feat: add CLIP 512d embedding to predict output

app.py

@@ -1,33 +1,102 @@
-import gradio as gr
-from ultralytics import YOLO
-import cv2
-
-# 모델 로드
-model = YOLO("best.pt")
-
-def predict(image):
-    # YOLO 추론
-    results = model(image)
-    
-    # 상의/하의 좌표 및 라벨 추출 (기존 프로젝트 JSON 형식에 맞춰 가공)
-    detections = []
-    for r in results:
-        for box in r.boxes:
-            x1, y1, x2, y2 = box.xyxy[0].tolist()
-            label = model.names[int(box.cls[0])]
-            conf = float(box.conf[0])
-            detections.append({
-                "label": label,
-                "confidence": conf,
-                "box": [int(x1), int(y1), int(x2), int(y2)]
-            })
-    return detections
-
-# Gradio 인터페이스 설정 (API 전용이므로 단순하게 구성)
-demo = gr.Interface(
-    fn=predict, 
-    inputs=gr.Image(type="numpy"), 
-    outputs="json"
-)
-
-demo.launch()
+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
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+logger.info("YOLO 모델 로딩 중...")
+yolo_model = YOLO("best.pt")
+logger.info("YOLO 로드 완료")
+
+logger.info("CLIP 모델 로딩 중 (openai/clip-vit-base-patch32)...")
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+clip_model.eval()
+logger.info("CLIP 로드 완료 dim=512")
+
+CATEGORY_MAP = {
+    "top": "top",
+    "bottom": "bottom",
+    "outer": "outer",
+}
+
+def _l2_normalize(vec):
+    norm = np.linalg.norm(vec)
+    if norm < 1e-9:
+        return vec.tolist()
+    return (vec / norm).tolist()
+
+def predict(image):
+    if image is None:
+        return {"boxes": [], "embedding": None, "label": "unknown", "category": None}
+
+    pil_img = Image.fromarray(image).convert("RGB")
+    w, h = pil_img.size
+
+    # 1. YOLO 의류 탐지
+    detections = []
+    try:
+        results = yolo_model(pil_img, conf=0.30, iou=0.6)
+        for r in results:
+            if r.boxes is None:
+                continue
+            for box in r.boxes:
+                x1, y1, x2, y2 = [int(v) for v in box.xyxy[0].tolist()]
+                conf = float(box.conf[0].item())
+                cls_id = int(box.cls[0].item())
+                raw_label = r.names.get(cls_id, "top")
+                x1, x2 = max(0, min(x1, w - 1)), max(x1 + 1, min(x2, w))
+                y1, y2 = max(0, min(y1, h - 1)), max(y1 + 1, min(y2, h))
+                detections.append({
+                    "label": raw_label,
+                    "confidence": round(conf, 4),
+                    "box": [x1, y1, x2, y2],
+                })
+        detections.sort(key=lambda b: b["confidence"], reverse=True)
+    except Exception as e:
+        logger.error(f"YOLO error: {e}")
+
+    # 2. 탐지 영역 크롭 (5% 패딩)
+    crop_img = pil_img
+    if detections:
+        b = detections[0]
+        x1, y1, x2, y2 = b["box"]
+        bw, bh = x2 - x1, y2 - y1
+        pad_w, pad_h = int(bw * 0.05), int(bh * 0.05)
+        cx1, cy1 = max(0, x1 - pad_w), max(0, y1 - pad_h)
+        cx2, cy2 = min(w, x2 + pad_w), min(h, y2 + pad_h)
+        if cx2 > cx1 and cy2 > cy1:
+            crop_img = pil_img.crop((cx1, cy1, cx2, cy2))
+
+    # 3. CLIP 512차원 임베딩 추출
+    embedding = None
+    try:
+        inputs = clip_processor(images=crop_img, return_tensors="pt")
+        with torch.no_grad():
+            features = clip_model.get_image_features(**inputs)
+        vec = features[0].cpu().numpy().astype(np.float32)
+        embedding = _l2_normalize(vec)
+        logger.info(f"CLIP embedding done dim={len(embedding)}")
+    except Exception as e:
+        logger.error(f"CLIP error: {e}")
+
+    top_label = detections[0]["label"] if detections else "unknown"
+    category = CATEGORY_MAP.get(top_label.strip().lower(), top_label.lower())
+
+    return {"boxes": detections, "embedding": embedding, "label": top_label, "category": category}
+
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy", label="Fashion Image"),
+    outputs=gr.JSON(label="Result (boxes + 512d embedding)"),
+    title="Lookalike YOLO + CLIP",
+    description="YOLOv11 clothing detection + CLIP 512d embedding extraction API",
+)
+
+demo.launch()