Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,178 @@
+"""
+P3 모델 -- Fashionpedia 2-Stage Model (Fashion Image Segmentation · Attribute Tagging)
+
+[Stage 1] build_maskrcnn() -- Mask R-CNN ResNet50-FPN v2 (COCO pretrained)
+    · box/mask predictor 를 Fashionpedia 47-class(46 카테고리 + background) 로 교체
+    · 입력: List[Tensor[3,H,W]] / 출력: train→loss dict, eval→예측 List[Dict]
+
+[Stage 2] AttributeClassifier -- ResNet-50 기반 multi-label 속성 분류
+    · 입력: Tensor[B,3,224,224] (배경 제거된 crop)
+    · 출력: logits [B,294]  (sigmoid 는 BCEWithLogitsLoss 내부에서)
+
+[헬퍼] build_model(stage, **kwargs) -- 'maskrcnn' / 'attribute' 분기
+
+[중요] Stage 1 label 은 dataset 에서 category_id+1 로 들어온다 (0=background).
+       따라서 num_classes=47, 추론 결과를 보여줄 땐 label-1 로 되돌린다.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+import torch.nn as nn
+
+import torchvision
+from torchvision.models.detection import (
+    maskrcnn_resnet50_fpn_v2,
+    MaskRCNN_ResNet50_FPN_V2_Weights,
+)
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
+
+
+# ============================================================================
+# Stage 1: Mask R-CNN
+# ============================================================================
+
+def build_maskrcnn(num_classes: int = 47, pretrained: bool = True) -> nn.Module:
+    """Mask R-CNN ResNet50-FPN v2 를 Fashionpedia 용으로 만든다.
+
+    COCO 로 사전학습된 모델을 불러온 뒤, 마지막 분류/마스크 head 만 우리 클래스 수
+    (47 = 46 카테고리 + background)에 맞게 갈아끼운다. backbone(특징 추출부)은
+    COCO 가중치를 그대로 살려 전이학습한다.
+
+    Args:
+        num_classes (int): 클래스 수 (background 포함). Fashionpedia 는 47.
+        pretrained (bool): COCO 사전학습 가중치 사용 여부.
+
+    Returns:
+        nn.Module: head 가 교체된 Mask R-CNN.
+            · train(): model(images, targets) → loss dict
+            · eval():  model(images)          → [{boxes,labels,scores,masks}, ...]
+    """
+    weights = MaskRCNN_ResNet50_FPN_V2_Weights.COCO_V1 if pretrained else None
+    model = maskrcnn_resnet50_fpn_v2(weights=weights)
+
+    # 1) 박스 분류 head 교체 (cls_score / bbox_pred 를 num_classes 로)
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+
+    # 2) 마스크 예측 head 교체
+    in_channels_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
+    hidden_dim = 256
+    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_channels_mask, hidden_dim, num_classes)
+
+    return model
+
+
+# ============================================================================
+# Stage 2: Attribute Classifier
+# ============================================================================
+
+class AttributeClassifier(nn.Module):
+    """ResNet-50 기반 multi-label 속성 분류기 (Stage 2).
+
+    배경이 제거된 인스턴스 crop 을 받아 294개 속성에 대한 logit 을 출력한다.
+    backbone(특징 추출부)은 ImageNet 사전학습 가중치를 쓰고, 그 위에 새 분류 head 를 얹는다.
+
+    구조:
+        backbone : ResNet-50 (GAP/fc 제거) → feature map [B, 2048, 7, 7]
+        head     : GAP → Dropout → Linear(2048→512) → ReLU → Dropout → Linear(512→294)
+
+    Args:
+        backbone (str): 'resnet50' (그 외 이름은 timm 으로 생성 시도).
+        num_attrs (int): 속성 수 (294).
+        pretrained (bool): ImageNet 사전학습 사용 여부.
+        dropout (float): 드롭아웃 비율.
+    """
+
+    def __init__(
+        self,
+        backbone: str = "resnet50",
+        num_attrs: int = 294,
+        pretrained: bool = True,
+        dropout: float = 0.3,
+    ) -> None:
+        super().__init__()
+        self.num_attrs = num_attrs
+
+        if backbone == "resnet50":
+            from torchvision.models import resnet50, ResNet50_Weights
+            weights = ResNet50_Weights.IMAGENET1K_V2 if pretrained else None
+            net = resnet50(weights=weights)
+            self.feature_dim = net.fc.in_features                 # 2048
+            # 마지막 avgpool(GAP) + fc 를 떼어내 feature map 추출부만 남긴다
+            self.backbone = nn.Sequential(*list(net.children())[:-2])
+        else:
+            import timm
+            # global_pool='' → GAP 미적용 feature map, num_classes=0 → fc 제거
+            self.backbone = timm.create_model(
+                backbone, pretrained=pretrained, num_classes=0, global_pool=""
+            )
+            self.feature_dim = self.backbone.num_features
+
+        # 분류 head (GAP 포함). sigmoid 는 적용하지 않음 (BCEWithLogitsLoss).
+        self.head = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Flatten(1),
+            nn.Dropout(dropout),
+            nn.Linear(self.feature_dim, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(512, num_attrs),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """순전파.
+
+        Args:
+            x (torch.Tensor): [B, 3, 224, 224]
+
+        Returns:
+            torch.Tensor: logits [B, num_attrs] (sigmoid 전)
+        """
+        features = self.backbone(x)        # [B, 2048, 7, 7]
+        logits = self.head(features)       # [B, num_attrs]
+        return logits
+
+    def get_param_groups(
+        self, backbone_lr: float = 1e-5, head_lr: float = 1e-4
+    ) -> List[dict]:
+        """backbone 과 head 에 서로 다른 학습률을 주는 파라미터 그룹.
+
+        사전학습된 backbone 은 살짝만(작은 lr), 새로 얹은 head 는 빠르게(큰 lr) 학습한다.
+
+        Args:
+            backbone_lr (float): backbone 학습률.
+            head_lr (float): head 학습률.
+
+        Returns:
+            List[dict]: optimizer 에 그대로 넘길 파라미터 그룹.
+        """
+        return [
+            {"params": self.backbone.parameters(), "lr": backbone_lr},
+            {"params": self.head.parameters(), "lr": head_lr},
+        ]
+
+
+# ============================================================================
+# 빌더 헬퍼
+# ============================================================================
+
+def build_model(stage: str, **kwargs) -> nn.Module:
+    """stage 이름으로 알맞은 모델을 만든다.
+
+    Args:
+        stage (str): 'maskrcnn'(Stage 1) 또는 'attribute'(Stage 2).
+        **kwargs: 각 빌더로 그대로 전달 (num_classes/pretrained 또는
+            backbone/num_attrs/pretrained/dropout 등).
+
+    Returns:
+        nn.Module: 생성된 모델.
+    """
+    if stage == "maskrcnn":
+        return build_maskrcnn(**kwargs)
+    if stage == "attribute":
+        return AttributeClassifier(**kwargs)
+    raise ValueError(f"알 수 없는 stage: {stage!r} (가능: 'maskrcnn', 'attribute')")