app.py

import gradio as gr
from ultralytics import YOLO
import numpy as np
from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPModel
import logging
import traceback

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

logger.info("YOLO 로딩...")
yolo_model = YOLO("best.pt")
logger.info("Fashion-CLIP 로딩...")
clip_model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")
clip_processor = CLIPProcessor.from_pretrained("patrickjohncyh/fashion-clip")
clip_model.eval()
logger.info("모델 로딩 완료!")

# ─── 탐지 품질 파라미터 ───────────────────────────────────────────────────────
# 신뢰도 임계값: 이보다 낮은 탐지 결과는 노이즈로 간주하여 제거
CONF_THRESHOLD = 0.35

# 최소 면적 비율: 이미지 전체 면적 대비 탐지 박스 면적이 이 비율 미만이면 제거
# (예: 0.04 = 전체 이미지의 4% 미만인 박스는 너무 작아 신뢰할 수 없음)
MIN_AREA_RATIO = 0.04

# IOU 임계값: YOLO 내부 NMS에서 중복 박스 제거 기준 (낮을수록 엄격)
IOU_THRESHOLD = 0.45


def _containment_ratio(inner: dict, outer: dict) -> float:
    """
    inner 박스가 outer 박스 내부에 포함된 비율(0.0~1.0)을 반환.
    1.0이면 inner가 outer에 완전히 포함됨.
    inner 면적 대비 교집합 면적의 비율로 계산.
    """
    ix1 = max(inner["x1"], outer["x1"])
    iy1 = max(inner["y1"], outer["y1"])
    ix2 = min(inner["x2"], outer["x2"])
    iy2 = min(inner["y2"], outer["y2"])

    inter_w = max(0.0, ix2 - ix1)
    inter_h = max(0.0, iy2 - iy1)
    inter_area = inter_w * inter_h

    inner_area = max(1.0, (inner["x2"] - inner["x1"]) * (inner["y2"] - inner["y1"]))
    return inter_area / inner_area


def _select_best_boxes(raw_boxes: list[dict], img_w: int, img_h: int) -> list[dict]:
    """
    YOLO가 반환한 모든 박스에서 품질 낮은 박스를 제거하고
    카테고리(레이블)별로 신뢰도가 가장 높은 박스 1개씩만 선택.

    추가 후처리:
    - 포함 관계 필터: 한 박스가 다른 박스 안에 크게 포함되면 제거
      예) Outer 박스 안에 Top 박스가 80% 이상 들어 있으면 Top 제거
    - Bottom 박스 확장: 하의 박스가 이미지 하단에 닿지 않으면 아래로 늘려 바지 전체 포함

    Args:
        raw_boxes: YOLO에서 반환된 원본 박스 목록
        img_w: 원본 이미지 가로 픽셀
        img_h: 원본 이미지 세로 픽셀

    Returns:
        정제된 박스 목록
    """
    img_area = img_w * img_h
    if img_area <= 0:
        return raw_boxes

    # 1단계: 신뢰도 + 최소 면적 기준으로 노이즈 박스 제거
    # Bottom 전용 임계값: 가려진 다리·오른쪽 바지 등이 낮은 신뢰도로 탐지될 수 있으므로
    # 다른 카테고리(0.35)보다 낮은 0.20을 적용해 2개 박스가 Union까지 살아남도록 함.
    BOTTOM_CONF_THRESHOLD = 0.20
    filtered = []
    for box in raw_boxes:
        label = box.get("label", "").lower()
        conf = box.get("confidence", 0.0)
        # 하의는 낮은 신뢰도 임계값 적용, 나머지는 기본 임계값
        threshold = BOTTOM_CONF_THRESHOLD if label in ("bottom", "하의") else CONF_THRESHOLD
        if conf < threshold:
            logger.debug(f"신뢰도 미달 박스 제거: label={label}, conf={conf:.3f} (기준={threshold:.2f})")
            continue

        x1, y1, x2, y2 = box["x1"], box["y1"], box["x2"], box["y2"]
        box_area = max(0.0, (x2 - x1)) * max(0.0, (y2 - y1))
        area_ratio = box_area / img_area

        if area_ratio < MIN_AREA_RATIO:
            logger.debug(
                f"면적 미달 박스 제거: label={label}, "
                f"area_ratio={area_ratio:.3f} (<{MIN_AREA_RATIO})"
            )
            continue

        filtered.append(box)

    # 2단계: 카테고리(레이블)별 박스를 Union으로 병합
    # 동일 카테고리(예: Bottom)가 여러 개 탐지된 경우,
    # 가장 신뢰도 높은 것만 고르지 않고 모든 박스를 합집합(Union)으로 합침.
    # → 바지가 두 박스로 나뉘어 탐지될 때 두 박스를 합쳐 전체 바지 영역 커버
    union_by_label: dict[str, dict] = {}
    for box in filtered:
        label = box.get("label", "unknown")
        if label not in union_by_label:
            # 첫 번째 박스는 그대로 복사 (원본 변경 방지)
            union_by_label[label] = dict(box)
        else:
            prev = union_by_label[label]
            # 기존 박스와 현재 박스의 합집합(Union) 계산
            prev["x1"] = min(prev["x1"], box["x1"])
            prev["y1"] = min(prev["y1"], box["y1"])
            prev["x2"] = max(prev["x2"], box["x2"])
            prev["y2"] = max(prev["y2"], box["y2"])
            # 신뢰도는 최대값 유지
            prev["confidence"] = max(prev.get("confidence", 0.0), box.get("confidence", 0.0))
            logger.info(
                f"박스 Union 병합: '{label}' 박스 2개 합산 "
                f"→ ({prev['x1']:.0f},{prev['y1']:.0f})-({prev['x2']:.0f},{prev['y2']:.0f})"
            )

    # 3단계: 포함 관계 필터
    # 박스 A가 박스 B 안에 CONTAINMENT_THRESHOLD 이상 포함되면 A를 제거
    # 예) Outer(큰 박스) 안에 Top(작은 박스)이 80%+ 포함 → Top 제거
    # 단, Bottom은 Outer 아래쪽에 별도 존재하므로 다른 기준 적용
    CONTAINMENT_THRESHOLD = 0.75  # inner 박스 면적의 이 비율 이상이 outer 안에 있으면 제거
    candidates = list(union_by_label.values())
    to_remove = set()

    for i, box_a in enumerate(candidates):
        for j, box_b in enumerate(candidates):
            if i == j:
                continue
            label_a = box_a.get("label", "").lower()
            label_b = box_b.get("label", "").lower()

            # bottom은 outer와 수직으로 분리되므로 포함 판단에서 제외
            if label_a in ("bottom", "하의") or label_b in ("bottom", "하의"):
                continue

            ratio = _containment_ratio(box_a, box_b)
            if ratio >= CONTAINMENT_THRESHOLD:
                # box_a가 box_b 안에 크게 포함됨 → box_a 면적이 box_b보다 작으면 제거
                area_a = (box_a["x2"] - box_a["x1"]) * (box_a["y2"] - box_a["y1"])
                area_b = (box_b["x2"] - box_b["x1"]) * (box_b["y2"] - box_b["y1"])
                if area_a < area_b:
                    to_remove.add(label_a)
                    logger.info(
                        f"포함 관계 필터: '{label_a}' 박스가 '{label_b}' 박스에 "
                        f"{ratio:.0%} 포함 → '{label_a}' 제거"
                    )

    result = [b for b in candidates if b.get("label", "").lower() not in to_remove]

    logger.info(
        f"박스 필터링: 원본 {len(raw_boxes)}개 → "
        f"신뢰도/면적 필터 후 {len(filtered)}개 → "
        f"Union 병합 후 {len(union_by_label)}개 → "
        f"포함 관계 필터 후 {len(result)}개"
    )
    return result



def _get_best_crop(pil_img: Image.Image, boxes: list[dict]) -> Image.Image:
    """
    필터링된 박스 중 신뢰도가 가장 높은 박스 영역을 크롭하여 반환.
    박스가 없으면 원본 이미지를 그대로 반환.

    크롭 이미지로 CLIP 임베딩을 생성하면
    전체 이미지 임베딩보다 패션 아이템에 집중된 더 정확한 벡터를 얻을 수 있음.
    """
    if not boxes:
        return pil_img

    # 신뢰도가 가장 높은 박스 선택
    best = max(boxes, key=lambda b: b.get("confidence", 0.0))
    x1 = int(best["x1"])
    y1 = int(best["y1"])
    x2 = int(best["x2"])
    y2 = int(best["y2"])

    # 원본 이미지 범위 클램핑
    w, h = pil_img.size
    x1 = max(0, min(x1, w - 1))
    x2 = max(x1 + 1, min(x2, w))
    y1 = max(0, min(y1, h - 1))
    y2 = max(y1 + 1, min(y2, h))

    try:
        cropped = pil_img.crop((x1, y1, x2, y2))
        logger.info(
            f"CLIP 임베딩용 크롭 이미지: "
            f"label={best.get('label')}, conf={best.get('confidence', 0):.3f}, "
            f"crop=({x1},{y1},{x2},{y2})"
        )
        return cropped
    except Exception as e:
        logger.warning(f"크롭 실패, 원본 사용: {e}")
        return pil_img


def predict(image):
    try:
        if image is None:
            return {
                "status": "error",
                "error_message": "No image provided",
                "embedding": None,
                "boxes": [],
                "label": "unknown",
                "category": None,
            }

        # ── 이미지 전처리 ─────────────────────────────────
        if isinstance(image, str):
            pil_img = Image.open(image).convert("RGB")
        elif isinstance(image, np.ndarray):
            pil_img = Image.fromarray(image).convert("RGB")
        elif isinstance(image, Image.Image):
            pil_img = image.convert("RGB")
        else:
            pil_img = Image.open(str(image)).convert("RGB")

        img_w, img_h = pil_img.size

        # ── YOLO 탐지 ─────────────────────────────────────
        # conf=0.10: 낮게 설정하여 가려진/약하게 탐지된 박스도 일단 수집
        #            실제 신뢰도 필터는 _select_best_boxes에서 카테고리별로 처리
        # iou=0.80: NMS를 느슨하게 → 겹치는 Bottom 박스 2개가 모두 살아남음
        #           (예: 왼쪽 다리 박스 + 오른쪽 다리 박스가 0.80 미만으로 겹침)
        #           두 박스가 모두 도달해야 Union으로 전체 바지 영역 합산 가능
        results = yolo_model.predict(
            source=pil_img,
            conf=0.10,       # 낮게: 약한 탐지도 수집 (이후 _select_best_boxes에서 필터)
            iou=0.80,        # 느슨한 NMS: 같은 카테고리 박스 2개가 모두 살아남도록
            save=False,
            verbose=False,
        )

        raw_boxes = []
        if results and len(results) > 0:
            for result in results:
                if result.boxes:
                    for box in result.boxes:
                        x1, y1, x2, y2 = box.xyxy[0].tolist()
                        conf = float(box.conf[0]) if box.conf is not None else 0
                        cls = int(box.cls[0]) if box.cls is not None else 0
                        label = (
                            result.names.get(cls, "unknown")
                            if hasattr(result, "names")
                            else "unknown"
                        )
                        raw_boxes.append(
                            {
                                "x1": x1,
                                "y1": y1,
                                "x2": x2,
                                "y2": y2,
                                "confidence": conf,
                                "label": label,
                            }
                        )

        # ── 박스 후처리: 노이즈 제거 + 카테고리별 최고 신뢰도 1개 선택 ────────
        filtered_boxes = _select_best_boxes(raw_boxes, img_w, img_h)

        # 대표 카테고리: 신뢰도 가장 높은 박스의 레이블
        detected_category = None
        if filtered_boxes:
            best_box = max(filtered_boxes, key=lambda b: b.get("confidence", 0.0))
            label = best_box.get("label", "")
            if label and label != "unknown":
                detected_category = label

        # ── CLIP 임베딩 (512d) ─────────────────────────────
        # 크롭 이미지로 임베딩 생성 (박스가 있으면 상품 영역만 크롭)
        embed_img = _get_best_crop(pil_img, filtered_boxes)
        inputs = clip_processor(images=embed_img, return_tensors="pt")

        with torch.no_grad():
            # vision_model → visual_projection 순서로 명시 호출
            vision_outputs = clip_model.vision_model(**inputs)
            features = clip_model.visual_projection(vision_outputs.pooler_output)

        # L2 정규화 (코사인 유사도 최적화)
        embedding = torch.nn.functional.normalize(features, p=2, dim=1)
        embedding_list = embedding[0].cpu().tolist()

        logger.info(
            f"임베딩 생성 완료: dim={len(embedding_list)}, "
            f"filtered_boxes={len(filtered_boxes)}, "
            f"category={detected_category}"
        )

        return {
            "status": "success",
            "embedding": embedding_list,   # 512d 벡터
            "boxes": filtered_boxes,
            "label": detected_category if detected_category else "full_image",
            "category": detected_category,
        }

    except Exception as e:
        err_msg = traceback.format_exc()
        logger.error(f"추론 중 예외 발생: {err_msg}")
        return {
            "status": "error",
            "error_message": str(e),
            "traceback": err_msg,
            "embedding": None,
            "boxes": [],
            "label": "unknown",
            "category": None,
        }


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs=gr.JSON(),
)
demo.launch(show_error=True)