import os
import torch
import numpy as np
import cv2

from PIL import Image

from transformers import (
    AutoImageProcessor,
    SegformerForSemanticSegmentation
)

# =========================================================
# 설정
# =========================================================

IMAGE_PATH = "test1.png"

MASK_OUTPUT = "top_mask.png"
PREVIEW_OUTPUT = "top_preview.png"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# =========================================================
# 모델 로드
# =========================================================

print("모델 로딩 중...")

processor = AutoImageProcessor.from_pretrained(
    "mattmdjaga/segformer_b2_clothes"
)

model = SegformerForSemanticSegmentation.from_pretrained(
    "mattmdjaga/segformer_b2_clothes"
)

model.to(DEVICE)
model.eval()

print("모델 로딩 완료")

# =========================================================
# 이미지 로드
# =========================================================

image = Image.open(IMAGE_PATH).convert("RGB")

width, height = image.size

# =========================================================
# 전처리
# =========================================================

inputs = processor(
    images=image,
    return_tensors="pt"
)

inputs = {
    k: v.to(DEVICE)
    for k, v in inputs.items()
}

# =========================================================
# 추론
# =========================================================

with torch.no_grad():

    outputs = model(**inputs)

# =========================================================
# 원본 크기로 복원
# =========================================================

upsampled_logits = torch.nn.functional.interpolate(
    outputs.logits,
    size=(height, width),
    mode="bilinear",
    align_corners=False
)

# =========================================================
# 클래스 예측
# =========================================================

pred_seg = upsampled_logits.argmax(dim=1)[0]

segmentation = pred_seg.cpu().numpy()

# =========================================================
# 클래스 확인
# =========================================================

unique_classes = np.unique(segmentation)

print("검출 클래스:", unique_classes)

# =========================================================
# 상의 클래스 ID
# =========================================================
#
# 모델마다 다를 수 있음
# 보통 upper-clothes 계열
#

TOP_IDS = [4]

# =========================================================
# 마스크 생성
# =========================================================

mask = np.isin(segmentation, TOP_IDS)

mask_uint8 = (mask.astype(np.uint8) * 255)

# =========================================================
# 저장
# =========================================================

mask_img = Image.fromarray(mask_uint8)

mask_img.save(MASK_OUTPUT)

print("마스크 저장 완료:", MASK_OUTPUT)

# =========================================================
# 미리보기 생성
# =========================================================

image_np = np.array(image)

preview = np.zeros((height, width, 4), dtype=np.uint8)

preview[..., :3] = image_np
preview[..., 3] = mask_uint8

preview_img = Image.fromarray(preview)

preview_img.save(PREVIEW_OUTPUT)

print("미리보기 저장 완료:", PREVIEW_OUTPUT)