p3-segmentation / utils.py
yeonjung77's picture
Upload utils.py with huggingface_hub
38ec06a verified
Raw
History Blame Contribute Delete
19.4 kB
"""
P3 (Fashion Image Segmentation · Attribute Tagging) -- 공통 유틸리티
2단계(2-Stage) 파이프라인 전반에서 공유하는 보조 도구를 모았습니다.
Stage 1 : Mask R-CNN instance segmentation (46개 카테고리)
Stage 2 : 인스턴스 crop 기반 multi-label attribute tagging (294개 속성)
[주요 구성]
1. 카테고리 상수 MAIN_GARMENT_NAMES(27) / GARMENT_PART_NAMES(19)
2. load_category_ids() 런타임에 COCO JSON → (MAIN_IDS, PART_IDS, id2name)
3. load_attribute_supers() 인덱스 → 속성 슈퍼카테고리 매핑 (그룹별 메트릭용)
4. set_seed() 재현성 시드 고정
5. AverageMeter 실행 평균 추적
6. EarlyStopping 조기 종료 (mAP@0.5 / F1 모니터링)
7. compute_segmentation_metrics() Stage 1 -- COCO mAP (Main/Part 분리)
8. compute_attribute_metrics() Stage 2 -- macro F1 + 슈퍼카테고리별 F1
9. collate_fn_stage1() Mask R-CNN용 variable-size collate
10. setup_logger() 파일 + 콘솔 동시 로깅
11. get_pos_weight() inverse-frequency BCE pos_weight
근거: 카테고리/속성 구조는 EDA(`004_EDA/P3_EDA_short.ipynb` cell-05)에서 검증.
"""
from __future__ import annotations
import os
import json
import random
import logging
from collections import Counter
from typing import Dict, List, Optional, Sequence, Tuple
import numpy as np
import torch
# ============================================================================
# 1. 카테고리 상수 (EDA 결과 기반 -- 004_EDA/P3_EDA_short.ipynb cell-05)
# ----------------------------------------------------------------------------
# Fashionpedia 46개 카테고리 = Main 의류 27개 + Part 부속/장식 19개.
# Main 만 Stage 2 속성 태깅의 입력으로 사용한다 (EDA 가정 ③: Main 속성이
# Part 보다 압도적으로 풍부). 이름은 COCO `categories[*].name` 과 정확히 일치.
# ============================================================================
#: Stage 2 입력이 되는 본체 의류 27개 (옷 + 액세서리)
MAIN_GARMENT_NAMES = {
# 상의
"shirt, blouse", "top, t-shirt, sweatshirt", "sweater", "cardigan",
"jacket", "vest",
# 하의 / 원피스류
"pants", "shorts", "skirt", "coat", "dress", "jumpsuit", "cape",
# 헤드 / 액세서리
"glasses", "hat", "headband, head covering, hair accessory",
"tie", "glove", "watch", "belt",
# 다리 / 발
"leg warmer", "tights, stockings", "sock", "shoe",
# 소품
"bag, wallet", "scarf", "umbrella",
}
#: 의류에 부속/장식되는 파트 19개 (Stage 2에서 제외, Stage 1에서만 검출)
GARMENT_PART_NAMES = {
# 구조 부위
"hood", "collar", "lapel", "epaulette", "sleeve", "pocket", "neckline",
"buckle", "zipper",
# 장식 요소
"applique", "bead", "bow", "flower", "fringe", "ribbon", "rivet",
"ruffle", "sequin", "tassel",
}
def load_category_ids(
coco_json_path: str,
) -> Tuple[set, set, Dict[int, str]]:
"""COCO JSON을 읽어 Main / Part 카테고리 ID 집합과 ID→이름 매핑을 만든다.
카테고리 ID 는 데이터셋 버전에 따라 달라질 수 있으므로 하드코딩하지 않고
런타임에 이름(MAIN_GARMENT_NAMES / GARMENT_PART_NAMES)으로 매핑한다.
Args:
coco_json_path (str): instances_attributes_*.json 경로
Returns:
Tuple[set, set, Dict[int, str]]:
- MAIN_IDS : Main 카테고리 ID 집합 (27개)
- PART_IDS : Part 카테고리 ID 집합 (19개)
- id2name : {category_id: name} 전체 매핑 (46개)
Raises:
ValueError: JSON 의 카테고리 이름이 정의된 상수와 어긋날 때
(Main+Part 합이 46이 아니거나 분류되지 않은 이름이 있을 때).
"""
with open(coco_json_path, "r") as f:
raw = json.load(f)
id2name = {c["id"]: c["name"] for c in raw["categories"]}
main_ids, part_ids, unknown = set(), set(), []
for cid, name in id2name.items():
if name in MAIN_GARMENT_NAMES:
main_ids.add(cid)
elif name in GARMENT_PART_NAMES:
part_ids.add(cid)
else:
unknown.append(name)
if unknown:
raise ValueError(
f"Main/Part 어느 쪽에도 속하지 않는 카테고리: {unknown}. "
"MAIN_GARMENT_NAMES / GARMENT_PART_NAMES 상수를 확인하세요."
)
if len(main_ids) != len(MAIN_GARMENT_NAMES) or len(part_ids) != len(GARMENT_PART_NAMES):
raise ValueError(
f"카테고리 개수 불일치: Main {len(main_ids)}/{len(MAIN_GARMENT_NAMES)}, "
f"Part {len(part_ids)}/{len(GARMENT_PART_NAMES)}"
)
return main_ids, part_ids, id2name
def load_attribute_supers(
coco_json_path: str,
num_attrs: int = 294,
) -> Tuple[List[str], Dict[int, int]]:
"""속성 인덱스 → 슈퍼카테고리 매핑을 만든다 (슈퍼카테고리별 F1 메트릭용).
Fashionpedia 속성 ID 는 0~340 범위에 비연속적으로 분포(총 294개)하므로,
학습용 multi-hot 벡터의 인덱스는 **속성 ID 오름차순 순서**로 0..293에 매핑한다.
preprocess 의 multi-hot 생성과 반드시 동일한 정렬 규칙을 사용해야 한다.
Args:
coco_json_path (str): instances_attributes_*.json 경로
num_attrs (int): 속성 개수 (기본 294)
Returns:
Tuple[List[str], Dict[int, int]]:
- idx2super : 길이 num_attrs 리스트. idx2super[i] = i번째 속성의 슈퍼카테고리
- attr_id2idx : {원본 attribute_id: 0..num_attrs-1 인덱스}
Note:
Fashionpedia 데이터의 실제 슈퍼카테고리는 11종이다
(animal, leather, length, neckline type, nickname, non-textile material type,
opening type, silhouette, textile finishing, textile pattern, waistline).
"""
with open(coco_json_path, "r") as f:
raw = json.load(f)
attrs = sorted(raw["attributes"], key=lambda a: a["id"])
if len(attrs) != num_attrs:
raise ValueError(f"속성 개수 불일치: JSON {len(attrs)} != num_attrs {num_attrs}")
idx2super = [a["supercategory"] for a in attrs]
attr_id2idx = {a["id"]: i for i, a in enumerate(attrs)}
return idx2super, attr_id2idx
# ============================================================================
# 2. 재현성
# ============================================================================
def set_seed(seed: int = 42) -> None:
"""random / numpy / torch / cudnn 시드를 모두 고정해 재현성을 확보한다.
Args:
seed (int): 랜덤 시드 (기본 42)
"""
os.environ["PYTHONHASHSEED"] = str(seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# 완전 결정론적 동작 (속도보다 재현성 우선)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# ============================================================================
# 3. 학습 보조 클래스
# ============================================================================
class AverageMeter:
"""실행 평균과 합계를 추적하는 누적기.
배치 단위 loss / 메트릭의 에폭 평균을 구할 때 사용한다.
사용 예:
meter = AverageMeter("loss")
for loss, bs in batches:
meter.update(loss, n=bs)
print(meter.avg)
"""
def __init__(self, name: str = "metric") -> None:
self.name = name
self.reset()
def reset(self) -> None:
"""모든 통계를 0으로 초기화."""
self.val = 0.0
self.avg = 0.0
self.sum = 0.0
self.count = 0
def update(self, val: float, n: int = 1) -> None:
"""새 값으로 통계를 갱신한다.
Args:
val (float): 현재 배치의 (평균) 값
n (int): 표본 수 (배치 크기). 가중 평균에 사용
"""
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count if self.count > 0 else 0.0
def __repr__(self) -> str:
return f"{self.name}: {self.avg:.4f} (n={self.count})"
class EarlyStopping:
"""검증 메트릭이 patience 에폭 동안 개선되지 않으면 학습 중단 신호를 준다.
Stage 1 은 mAP@0.5 (mode='max'), Stage 2 는 macro F1 (mode='max') 을 모니터링.
loss 처럼 작아야 좋은 지표는 mode='min' 으로 사용.
"""
def __init__(
self,
patience: int = 5,
mode: str = "max",
delta: float = 0.0,
) -> None:
"""초기화.
Args:
patience (int): 개선 없이 허용할 최대 에폭 수
mode (str): "max"(클수록 좋음) 또는 "min"(작을수록 좋음)
delta (float): 개선으로 인정할 최소 변화량
"""
if mode not in ("max", "min"):
raise ValueError(f"mode 는 'max' 또는 'min' 이어야 합니다 (받은 값: {mode})")
self.patience = patience
self.mode = mode
self.delta = delta
self.counter = 0
self.best_score: Optional[float] = None
self.early_stop = False
def _is_improved(self, score: float) -> bool:
if self.best_score is None:
return True
if self.mode == "max":
return score > self.best_score + self.delta
return score < self.best_score - self.delta
def __call__(self, metric: float) -> bool:
"""현재 에폭 메트릭을 받아 학습 중단 여부를 반환한다.
Args:
metric (float): 현재 에폭의 검증 메트릭
Returns:
bool: True 이면 학습을 멈춰야 함
"""
if self._is_improved(metric):
self.best_score = metric
self.counter = 0
else:
self.counter += 1
if self.counter >= self.patience:
self.early_stop = True
return self.early_stop
# ============================================================================
# 4. collate_fn (Stage 1)
# ============================================================================
def collate_fn_stage1(batch: List[Tuple]) -> Tuple[List, List]:
"""Mask R-CNN용 collate 함수.
torchvision detection 모델은 서로 다른 크기의 이미지를 하나의 텐서로
stack 하지 않고 **리스트** 그대로 입력받는다. 기본 collate 의 자동 stacking
을 우회하기 위해 (이미지 리스트, 타겟 리스트) 로만 분리한다.
Args:
batch (List[Tuple]): [(image_tensor, target_dict), ...]
Returns:
Tuple[List, List]: (images, targets) -- 둘 다 길이 B 리스트
"""
images, targets = zip(*batch)
return list(images), list(targets)
# ============================================================================
# 5. Stage 1 메트릭 -- COCO mAP
# ============================================================================
def compute_segmentation_metrics(
coco_gt,
coco_dt,
main_ids: Optional[set] = None,
part_ids: Optional[set] = None,
iou_type: str = "segm",
) -> Dict[str, float]:
"""pycocotools COCOeval 로 instance segmentation mAP 를 계산한다.
Args:
coco_gt: GT COCO 객체 (`pycocotools.coco.COCO`)
coco_dt: 예측 결과 COCO 객체. 보통 `coco_gt.loadRes(results)` 로 생성하며,
results 는 [{"image_id", "category_id", "segmentation"(RLE),
"score"}, ...] 형식의 리스트.
main_ids (set, optional): Main 카테고리 ID 집합. 주면 Main mAP 분리 계산.
part_ids (set, optional): Part 카테고리 ID 집합. 주면 Part mAP 분리 계산.
iou_type (str): "segm"(마스크) 또는 "bbox".
Returns:
Dict[str, float]: {
"mAP_50": mAP @ IoU=0.50,
"mAP_50_95": mAP @ IoU=0.50:0.95 (COCO primary),
"main_mAP_50": (main_ids 제공 시) Main 카테고리 mAP@0.5,
"part_mAP_50": (part_ids 제공 시) Part 카테고리 mAP@0.5,
}
Note:
COCOeval.stats[0] = mAP@[.5:.95], stats[1] = mAP@0.5.
값이 -1 이면 해당 카테고리에 대한 예측/GT 가 없다는 뜻.
"""
from pycocotools.cocoeval import COCOeval
def _run(cat_ids: Optional[Sequence[int]] = None) -> Tuple[float, float]:
evaluator = COCOeval(coco_gt, coco_dt, iouType=iou_type)
if cat_ids is not None:
evaluator.params.catIds = list(cat_ids)
evaluator.evaluate()
evaluator.accumulate()
evaluator.summarize()
# stats[0]=mAP@[.5:.95], stats[1]=mAP@0.5
return float(evaluator.stats[1]), float(evaluator.stats[0])
map_50, map_50_95 = _run()
metrics: Dict[str, float] = {"mAP_50": map_50, "mAP_50_95": map_50_95}
if main_ids:
metrics["main_mAP_50"] = _run(main_ids)[0]
if part_ids:
metrics["part_mAP_50"] = _run(part_ids)[0]
return metrics
# ============================================================================
# 6. Stage 2 메트릭 -- multi-label attribute F1
# ============================================================================
def compute_attribute_metrics(
logits,
targets,
threshold: float = 0.5,
attr_supers: Optional[Sequence[str]] = None,
) -> Dict[str, float]:
"""multi-label 속성 분류 메트릭 (macro F1/precision/recall + 슈퍼카테고리별 F1).
Args:
logits: 모델 출력 [N, num_attrs]. **raw logit** 으로 보고 내부에서
sigmoid 를 적용한다 (torch.Tensor / np.ndarray 모두 허용).
targets: multi-hot 정답 [N, num_attrs] (0/1).
threshold (float): sigmoid 확률을 1로 판정할 임계값.
attr_supers (Sequence[str], optional): 길이 num_attrs. 각 속성 인덱스의
슈퍼카테고리 이름 (load_attribute_supers 로 얻음). 주면 그룹별
macro F1 을 함께 반환한다.
Returns:
Dict[str, float]: {
"f1_macro", "precision_macro", "recall_macro",
"f1_micro",
"f1_super/<supercategory>": ... (attr_supers 제공 시 그룹별)
}
"""
from sklearn.metrics import f1_score, precision_score, recall_score
probs = _to_numpy(logits)
y_true = _to_numpy(targets).astype(int)
y_pred = (_sigmoid(probs) >= threshold).astype(int)
metrics = {
"f1_macro": float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
"precision_macro": float(precision_score(y_true, y_pred, average="macro", zero_division=0)),
"recall_macro": float(recall_score(y_true, y_pred, average="macro", zero_division=0)),
"f1_micro": float(f1_score(y_true, y_pred, average="micro", zero_division=0)),
}
# 슈퍼카테고리(9~11개 그룹)별 macro F1
if attr_supers is not None:
if len(attr_supers) != y_true.shape[1]:
raise ValueError(
f"attr_supers 길이({len(attr_supers)})가 속성 차원({y_true.shape[1]})과 다릅니다."
)
attr_supers = np.asarray(attr_supers)
for group in sorted(set(attr_supers.tolist())):
cols = np.where(attr_supers == group)[0]
f1_g = f1_score(
y_true[:, cols], y_pred[:, cols], average="macro", zero_division=0
)
metrics[f"f1_super/{group}"] = float(f1_g)
return metrics
# ============================================================================
# 7. 로깅
# ============================================================================
def setup_logger(log_dir: str, name: str = "train") -> logging.Logger:
"""파일과 콘솔에 동시 출력하는 로거를 만든다.
Args:
log_dir (str): 로그 파일을 저장할 디렉토리 (없으면 생성)
name (str): 로거 이름 겸 로그 파일명 (`<name>.log`)
Returns:
logging.Logger: 설정된 로거
"""
os.makedirs(log_dir, exist_ok=True)
logger = logging.getLogger(name)
logger.setLevel(logging.INFO)
# 중복 핸들러 방지 (재호출 시 로그가 여러 번 찍히는 문제)
if logger.handlers:
return logger
logger.propagate = False
fmt = logging.Formatter(
"[%(asctime)s] %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
)
file_handler = logging.FileHandler(os.path.join(log_dir, f"{name}.log"))
file_handler.setFormatter(fmt)
logger.addHandler(file_handler)
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(fmt)
logger.addHandler(stream_handler)
return logger
# ============================================================================
# 8. BCE pos_weight (속성 long-tail 보정 -- EDA 가정 ⑤)
# ============================================================================
def get_pos_weight(
attr_counter: Counter,
num_attrs: int = 294,
max_weight: float = 10.0,
attr_ids: Optional[Sequence[int]] = None,
) -> torch.Tensor:
"""속성 빈도의 역수로 BCEWithLogitsLoss 의 pos_weight 를 계산한다.
희귀 속성(long-tail)일수록 양성 신호가 약하므로 손실에서 더 크게 가중한다.
가장 흔한 속성을 기준(1.0)으로 한 inverse-frequency 가중을 사용하고,
극단적 가중을 막기 위해 max_weight 로 clipping 한다.
pos_weight[i] = clip(max(count) / count[i], 1.0, max_weight)
Args:
attr_counter (Counter): {속성키: 등장횟수}. 키는 attr_ids 와 같은 체계여야 함.
num_attrs (int): 속성 개수 (출력 벡터 길이, 기본 294)
max_weight (float): pos_weight 상한 (기본 10.0)
attr_ids (Sequence[int], optional): 인덱스 i 에 대응하는 원본 속성 ID 목록.
Fashionpedia 처럼 ID 가 비연속(0~340)일 때 정렬된 ID 리스트를 넘긴다.
None 이면 키를 0..num_attrs-1 인덱스로 간주(연속 가정).
Returns:
torch.Tensor: shape (num_attrs,), dtype float32.
"""
keys = list(attr_ids) if attr_ids is not None else list(range(num_attrs))
if len(keys) != num_attrs:
raise ValueError(f"attr_ids 길이({len(keys)})가 num_attrs({num_attrs})와 다릅니다.")
# 등장 0회는 1로 바닥 처리 (0 나눗셈 방지)
counts = np.array([max(attr_counter.get(k, 0), 1) for k in keys], dtype=np.float64)
pos_weight = counts.max() / counts
pos_weight = np.clip(pos_weight, 1.0, max_weight)
return torch.tensor(pos_weight, dtype=torch.float32)
# ============================================================================
# 내부 헬퍼
# ============================================================================
def _to_numpy(x) -> np.ndarray:
"""torch.Tensor / np.ndarray / list 를 numpy 배열로 변환."""
if isinstance(x, torch.Tensor):
return x.detach().cpu().numpy()
return np.asarray(x)
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)))