model.py

"""
Faster R-CNN 밑바닥 구현 — [3/5] 모델
==========================================
전체 구조:

  이미지
    │  [백본: ResNet50 (ImageNet 사전학습)]
    ▼
  특징맵 (stride 16)
    │
    ├─▶ [RPN] ── 앵커마다 (객체 여부, 박스 보정) 예측 ── 후보영역(proposal) 생성
    │
    ▼
  [RoI Align] ── 각 후보영역을 고정 크기(7x7) 특징으로 추출
    │
    ▼
  [RoI Head] ── (클래스 분류, 클래스별 박스 보정) 최종 예측

2-stage 탐지기의 정석 구조. RPN이 "어디에 뭔가 있다"를,
RoI Head가 "그게 무엇이고 정확한 위치는 어디"를 담당한다.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision.ops import roi_align

from box_utils import (generate_anchors, shift_anchors, box_iou,
                       encode_boxes, decode_boxes, clip_boxes, nms)


# ---------------------------------------------------------------
# 백본: ResNet50의 conv1~layer3까지 (stride 16 특징맵)
# ---------------------------------------------------------------
class Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        resnet = torchvision.models.resnet50(weights="IMAGENET1K_V1")
        # layer4는 RoI Head에서 쓰고, 여기선 layer3까지 → stride 16
        self.body = nn.Sequential(
            resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
            resnet.layer1, resnet.layer2, resnet.layer3,
        )
        self.out_channels = 1024  # layer3 출력 채널
        self.stride = 16

        # 초기 레이어는 동결(작은 데이터셋 과적합·불안정 방지)
        for p in self.body[:5].parameters():
            p.requires_grad = False

    def forward(self, x):
        return self.body(x)


# ---------------------------------------------------------------
# RPN (Region Proposal Network)
# ---------------------------------------------------------------
class RPN(nn.Module):
    def __init__(self, in_channels=1024, num_anchors=9):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, 512, 3, padding=1)
        # 앵커마다: 객체/배경 2점수 → 여기선 objectness 1개(logit)로 단순화
        self.cls = nn.Conv2d(512, num_anchors, 1)          # objectness
        self.reg = nn.Conv2d(512, num_anchors * 4, 1)       # 박스 보정
        self.num_anchors = num_anchors

        for layer in [self.conv, self.cls, self.reg]:
            nn.init.normal_(layer.weight, std=0.01)
            nn.init.constant_(layer.bias, 0)

    def forward(self, feat):
        t = F.relu(self.conv(feat))
        logits = self.cls(t)   # [B, A, H, W]
        deltas = self.reg(t)   # [B, A*4, H, W]

        B, _, H, W = logits.shape
        # [B, H*W*A] 형태로 정리
        logits = logits.permute(0, 2, 3, 1).reshape(B, -1)
        deltas = deltas.permute(0, 2, 3, 1).reshape(B, -1, 4)
        return logits, deltas, (H, W)


# ---------------------------------------------------------------
# RoI Head (분류 + 박스 회귀)
# ---------------------------------------------------------------
class RoIHead(nn.Module):
    def __init__(self, num_classes, in_channels=1024, roi_size=7):
        super().__init__()
        # layer4 대신 간단한 FC 헤드로 구성(학습 가벼움)
        self.roi_size = roi_size
        flat = in_channels * roi_size * roi_size
        self.fc = nn.Sequential(
            nn.Linear(flat, 1024), nn.ReLU(inplace=True),
            nn.Linear(1024, 1024), nn.ReLU(inplace=True),
        )
        self.cls = nn.Linear(1024, num_classes)             # 클래스 분류
        self.reg = nn.Linear(1024, num_classes * 4)          # 클래스별 박스 보정
        self.num_classes = num_classes

        nn.init.normal_(self.cls.weight, std=0.01)
        nn.init.normal_(self.reg.weight, std=0.001)
        nn.init.constant_(self.cls.bias, 0)
        nn.init.constant_(self.reg.bias, 0)

    def forward(self, feat, rois, stride):
        """
        feat  : 백본 특징맵 [B,C,H,W]
        rois  : 후보영역 [N,4] (이미지 좌표)
        """
        # roi_align: 각 후보영역을 7x7 고정 크기 특징으로 추출
        # batch index 0 (batch_size=1 가정) 를 앞에 붙인다
        batch_idx = torch.zeros((rois.shape[0], 1), device=rois.device)
        rois_b = torch.cat([batch_idx, rois], dim=1)  # [N,5]
        pooled = roi_align(feat, rois_b,
                           output_size=(self.roi_size, self.roi_size),
                           spatial_scale=1.0 / stride,
                           sampling_ratio=2)
        x = pooled.flatten(1)
        x = self.fc(x)
        return self.cls(x), self.reg(x)


# ---------------------------------------------------------------
# 전체 모델
# ---------------------------------------------------------------
class FasterRCNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.backbone = Backbone()
        self.rpn = RPN(self.backbone.out_channels)
        self.head = RoIHead(num_classes, self.backbone.out_channels)
        self.num_classes = num_classes

        self.base_anchors = generate_anchors()  # 9개
        self.stride = self.backbone.stride

        # 학습/추론 하이퍼파라미터
        self.rpn_pre_nms = 12000
        self.rpn_post_nms_train = 2000
        self.rpn_post_nms_test = 300
        self.rpn_nms_thresh = 0.7

    def _anchors_for(self, feat_h, feat_w, device):
        a = shift_anchors(self.base_anchors, feat_h, feat_w, self.stride)
        return a.to(device)

    def _proposals(self, rpn_logits, rpn_deltas, anchors, img_hw, training):
        """RPN 출력 → NMS 거친 후보영역(proposal) 생성."""
        scores = torch.sigmoid(rpn_logits[0])         # [K]
        deltas = rpn_deltas[0]                          # [K,4]
        proposals = decode_boxes(deltas, anchors)
        proposals = clip_boxes(proposals, img_hw[0], img_hw[1])

        # 점수 상위만 추린 뒤 NMS
        n_pre = min(self.rpn_pre_nms, scores.numel())
        top = scores.topk(n_pre).indices
        proposals, scores = proposals[top], scores[top]

        keep = nms(proposals, scores, self.rpn_nms_thresh)
        n_post = self.rpn_post_nms_train if training else self.rpn_post_nms_test
        keep = keep[:n_post]
        return proposals[keep], scores[keep]

    def forward(self, image, target=None):
        """
        image  : [1,3,H,W] 단일 이미지(batch_size=1)
        target : 학습 시 {'boxes','labels'}, 추론 시 None
        """
        device = image.device
        img_h, img_w = image.shape[-2:]
        feat = self.backbone(image)
        _, _, fh, fw = feat.shape

        anchors = self._anchors_for(fh, fw, device)
        rpn_logits, rpn_deltas, _ = self.rpn(feat)

        proposals, _ = self._proposals(
            rpn_logits, rpn_deltas, anchors, (img_h, img_w),
            training=self.training)

        if self.training:
            # 학습 경로: 손실 계산 (train.py 의 헬퍼가 담당)
            return {
                "feat": feat, "anchors": anchors,
                "rpn_logits": rpn_logits[0], "rpn_deltas": rpn_deltas[0],
                "proposals": proposals, "stride": self.stride,
            }
        else:
            # 추론 경로: 최종 탐지 결과
            cls_logits, reg = self.head(feat, proposals, self.stride)
            return self._postprocess(cls_logits, reg, proposals, (img_h, img_w))

    @torch.no_grad()
    def _postprocess(self, cls_logits, reg, proposals, img_hw,
                     score_thresh=0.05, nms_thresh=0.3, max_det=100):
        """RoI Head 출력 → 클래스별 NMS → 최종 박스/라벨/점수."""
        probs = F.softmax(cls_logits, dim=1)       # [N, C]
        N, C = probs.shape
        reg = reg.reshape(N, C, 4)

        all_boxes, all_scores, all_labels = [], [], []
        for c in range(1, C):  # 0=배경 제외
            scores_c = probs[:, c]
            mask = scores_c > score_thresh
            if mask.sum() == 0:
                continue
            boxes_c = decode_boxes(reg[mask, c], proposals[mask])
            boxes_c = clip_boxes(boxes_c, img_hw[0], img_hw[1])
            scores_c = scores_c[mask]

            keep = nms(boxes_c, scores_c, nms_thresh)
            all_boxes.append(boxes_c[keep])
            all_scores.append(scores_c[keep])
            all_labels.append(torch.full((keep.numel(),), c,
                                         dtype=torch.int64, device=boxes_c.device))

        if not all_boxes:
            return {"boxes": torch.empty((0, 4)), "labels": torch.empty((0,), dtype=torch.int64),
                    "scores": torch.empty((0,))}

        boxes = torch.cat(all_boxes)
        scores = torch.cat(all_scores)
        labels = torch.cat(all_labels)

        # 전체에서 점수 상위 max_det개만
        if scores.numel() > max_det:
            top = scores.topk(max_det).indices
            boxes, scores, labels = boxes[top], scores[top], labels[top]
        return {"boxes": boxes, "labels": labels, "scores": scores}