postprocess.py

"""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)))