Upload postprocess.py with huggingface_hub

postprocess.py

@@ -0,0 +1,190 @@
+"""P3 후처리 -- Stage 1/2 결과 통합 + 박스 기반 Part 귀속 (Phase 4)
+
+Stage 1(Mask R-CNN)이 뱉는 평면 인스턴스 리스트를 "본체(Main) 의류 + 그에 달린
+부위(Part)" 트리로 재구성한다. 핵심은 **단순 이미지 공존이 아니라 박스 IoU/포함도로
+Part 를 가장 알맞은 Main 에 귀속**시키는 것 -- EDA 에서 본 부자연스러운 매핑
+(pants→sleeve, shoe→sleeve 등 같은 사진에 있다는 이유만의 연결)을 차단한다.
+
+구성:
+    box_iou(a, b)                 -- 두 박스 [x1,y1,x2,y2] IoU
+    box_containment(part, main)   -- Part 가 Main 에 포함된 비율 (intersection / part_area)
+    group_instances(...)          -- 평면 리스트 → (mains_with_parts, orphan_parts)
+    build_response(...)           -- 트리 + Stage 2 logits → 응답 JSON
+
+[detection 형식] 각 인스턴스는 dict:
+    {"category_id": int(0~45 원본), "category": str, "box": [x1,y1,x2,y2], "score": float}
+    Main 인스턴스에는 pipeline 이 "attr_index"(Stage 2 logits 행 번호)를 달아 둔다.
+
+- utils.py / model.py / dataset.py 등과 양식 통일.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+
+# ============================================================================
+# 1. 기하 -- 박스 IoU / 포함도
+# ============================================================================
+
+def box_iou(box_a: Sequence[float], box_b: Sequence[float]) -> float:
+    """두 박스 [x1, y1, x2, y2] 의 IoU (교집합 / 합집합).
+
+    Args:
+        box_a, box_b: [x1, y1, x2, y2] 형식 박스.
+
+    Returns:
+        float: IoU (0.0 ~ 1.0). 겹치지 않으면 0.0.
+    """
+    inter = _intersection_area(box_a, box_b)
+    if inter <= 0:
+        return 0.0
+    area_a = max(0.0, box_a[2] - box_a[0]) * max(0.0, box_a[3] - box_a[1])
+    area_b = max(0.0, box_b[2] - box_b[0]) * max(0.0, box_b[3] - box_b[1])
+    union = area_a + area_b - inter
+    return inter / union if union > 0 else 0.0
+
+
+def box_containment(part_box: Sequence[float], main_box: Sequence[float]) -> float:
+    """Part 박스가 Main 박스에 얼마나 들어가 있나 = 교집합 / Part 넓이.
+
+    IoU 와 달리 크기 차이에 둔감해, "작은 소매가 큰 상의 안에 거의 들어 있다"
+    같은 포함 관계를 잘 잡는다.
+
+    Args:
+        part_box: 포함되는 쪽(작은) 박스 [x1,y1,x2,y2].
+        main_box: 포함하는 쪽(큰) 박스 [x1,y1,x2,y2].
+
+    Returns:
+        float: 포함 비율 (0.0 ~ 1.0). Part 넓이가 0이면 0.0.
+    """
+    inter = _intersection_area(part_box, main_box)
+    part_area = max(0.0, part_box[2] - part_box[0]) * max(0.0, part_box[3] - part_box[1])
+    return inter / part_area if part_area > 0 else 0.0
+
+
+def _intersection_area(box_a: Sequence[float], box_b: Sequence[float]) -> float:
+    """두 박스의 교집합 넓이."""
+    x1 = max(box_a[0], box_b[0])
+    y1 = max(box_a[1], box_b[1])
+    x2 = min(box_a[2], box_b[2])
+    y2 = min(box_a[3], box_b[3])
+    return max(0.0, x2 - x1) * max(0.0, y2 - y1)
+
+
+# ============================================================================
+# 2. 인스턴스 그룹화 -- 평면 리스트 → Main 트리 + orphan
+# ============================================================================
+
+def group_instances(
+    detections: List[Dict],
+    main_ids: set,
+    part_ids: set,
+    iou_threshold: float = 0.1,
+    contain_threshold: float = 0.5,
+) -> Tuple[List[Dict], List[Dict]]:
+    """평면 detection 리스트를 Main(부위 포함) 트리로 재구성한다.
+
+    각 Part 는 **IoU 와 포함도가 가장 큰 Main** 에 귀속된다. 단, 그 Main 과의
+    관계가 `iou >= iou_threshold` 또는 `containment >= contain_threshold` 중
+    하나라도 만족할 때만 귀속하고, 둘 다 아니면 orphan(고아 Part)으로 둔다.
+
+    Args:
+        detections: Stage 1 인스턴스 dict 리스트 ({category_id, box, score, ...}).
+        main_ids: Main 카테고리 ID 집합 (utils.load_category_ids).
+        part_ids: Part 카테고리 ID 집합.
+        iou_threshold: 귀속 판정 IoU 하한.
+        contain_threshold: 귀속 판정 포함도 하한.
+
+    Returns:
+        Tuple[List[Dict], List[Dict]]:
+            - mains_with_parts: Main dict 리스트. 각 dict 에 "parts": [Part dict, ...] 추가.
+            - orphan_parts: 어느 Main 에도 귀속되지 못한 Part dict 리스트.
+    """
+    mains = [dict(d, parts=[]) for d in detections if d["category_id"] in main_ids]
+    parts = [d for d in detections if d["category_id"] in part_ids]
+
+    orphan_parts: List[Dict] = []
+    for part in parts:
+        best_i, best_iou, best_cont, best_key = -1, 0.0, 0.0, -1.0
+        for i, main in enumerate(mains):
+            iou = box_iou(part["box"], main["box"])
+            cont = box_containment(part["box"], main["box"])
+            key = max(iou, cont)            # IoU·포함도 중 큰 값으로 ���적 Main 선택
+            if key > best_key:
+                best_i, best_iou, best_cont, best_key = i, iou, cont, key
+
+        # 최적 Main 과의 관계가 임계값을 하나라도 넘으면 귀속, 아니면 orphan
+        if best_i >= 0 and (best_iou >= iou_threshold or best_cont >= contain_threshold):
+            mains[best_i]["parts"].append(part)
+        else:
+            orphan_parts.append(part)
+
+    return mains, orphan_parts
+
+
+# ============================================================================
+# 3. 응답 빌더 -- 트리 + Stage 2 속성 → JSON
+# ============================================================================
+
+def build_response(
+    mains_with_parts: List[Dict],
+    orphan_parts: List[Dict],
+    attribute_logits: Optional[np.ndarray],
+    threshold: float = 0.5,
+    id2attr_name: Optional[Sequence[str]] = None,
+) -> Dict:
+    """Main 트리 + Stage 2 logits 를 최종 응답 JSON 으로 만든다.
+
+    각 Main 의 속성은 `attribute_logits[main["attr_index"]]` 를 sigmoid 한 뒤
+    threshold 를 넘는 속성만 모아 이름 리스트로 변환한다.
+
+    Args:
+        mains_with_parts: group_instances 의 첫 반환값 (각 Main 에 "attr_index" 포함).
+        orphan_parts: group_instances 의 둘째 반환값.
+        attribute_logits: Stage 2 raw logits [M, num_attrs] (M=Main 수). 없으면 속성 빈칸.
+        threshold: 속성 판정 임계값 (sigmoid 확률 기준).
+        id2attr_name: 속성 인덱스→이름. 없으면 "attr_<idx>" 로 표기.
+
+    Returns:
+        Dict: {"garments": [...], "orphan_parts": [...]}
+    """
+    logits = None if attribute_logits is None else np.asarray(attribute_logits)
+
+    garments = []
+    for main in mains_with_parts:
+        attrs: List[str] = []
+        idx = main.get("attr_index")
+        if logits is not None and idx is not None and 0 <= idx < len(logits):
+            probs = _sigmoid(logits[idx])
+            for a in np.where(probs > threshold)[0]:
+                attrs.append(id2attr_name[a] if id2attr_name is not None else f"attr_{a}")
+
+        garments.append({
+            "category": main.get("category", str(main["category_id"])),
+            "box": [float(v) for v in main["box"]],
+            "score": float(main["score"]),
+            "attributes": attrs,
+            "parts": [_part_dict(p) for p in main["parts"]],
+        })
+
+    return {
+        "garments": garments,
+        "orphan_parts": [_part_dict(p) for p in orphan_parts],
+    }
+
+
+def _part_dict(p: Dict) -> Dict:
+    """Part 인스턴스를 응답용 최소 dict 로 정리."""
+    return {
+        "category": p.get("category", str(p["category_id"])),
+        "box": [float(v) for v in p["box"]],
+        "score": float(p["score"]),
+    }
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    """수치 안정 sigmoid."""
+    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))