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