"""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_" 로 표기. 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)))