src/detection_net.py

import torch
import torch.nn as nn
import torchvision.ops as ops
from torchvision.models import resnet50, ResNet50_Weights

class DetectionHead(nn.Module):
    def __init__(self, num_classes=20, pretrained=True):
        super().__init__()

        self.num_classes = num_classes

        weights = ResNet50_Weights.DEFAULT if pretrained else None
        self.conv5_x = resnet50(weights=weights).layer4

        self.fc_cls = nn.Linear(2048, num_classes + 1)
        self.fc_reg = nn.Linear(2048, num_classes * 4)

        self.fc_cls.weight.data.normal_(0, 0.01)
        self.fc_reg.weight.data.normal_(0, 0.001)
        self.fc_cls.bias.data.zero_()
        self.fc_reg.bias.data.zero_()

        for module in self.conv5_x.modules():
            if isinstance(module, nn.BatchNorm2d):
                module.weight.requires_grad = False
                module.bias.requires_grad = False

    def train(self, mode=True):
            super().train(mode)
    
            # requires_grad=False only freezes the affine params
            # training mode may again put the batchnorm layers in training mode, which is not what we want
            for module in self.conv5_x.modules():
                if isinstance(module, nn.BatchNorm2d):
                    module.eval()
    
            return self

    def forward(self, pooled_batch):
        # pooled_batch: list of B tensors [N_i, 1024, 7, 7] from RoIPool
        counts = [pooled.shape[0] for pooled in pooled_batch]

        # No proposals survived anywhere in the batch -- torch.cat would still work, but
        # conv5_x on a 0-length batch is pointless, so return correctly shaped empties
        if sum(counts) == 0:
            cls_logits = [pooled.new_zeros((0, self.num_classes + 1)) for pooled in pooled_batch]
            box_deltas = [pooled.new_zeros((0, self.num_classes * 4)) for pooled in pooled_batch]
            return cls_logits, box_deltas

        # One batched pass over every RoI in the batch, then split back per image
        x = torch.cat(pooled_batch, dim=0)           # [sumN, 1024, 7, 7]
        x = self.conv5_x(x)                          # [sumN, 2048, 4, 4]
        x = x.mean(dim=(2, 3))                       # global average pool -> [sumN, 2048]

        cls_logits = self.fc_cls(x)                  # [sumN, 21] raw logits
        box_deltas = self.fc_reg(x)                  # [sumN, 80] class-specific deltas

        return list(cls_logits.split(counts)), list(box_deltas.split(counts))


class DetectionLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, feature_maps, batch_proposals, batch_gt_boxes, batch_gt_labels,
                batch_img_height, batch_img_width, roi_pool, detection_head,
                num_samples=64, pos_fraction=0.25, background_label=20):
        # feature_maps:   [B, 1024, H_f, W_f] from Backbone
        # batch_proposals: list of B x [N_i, 4] CORNER absolute px, from RegionProposalNetwork
        # batch_gt_boxes:  list of B x [M_i, 4] CORNER absolute px
        # batch_gt_labels: list of B x [M_i] in 0..19
        device = feature_maps.device
        batch_size = len(batch_proposals)

        batch_sampled_proposals = []
        batch_labels = []
        batch_reg_targets = []

        # 64 sampled RoIs per image instead of all ~2000 proposals.
        for i in range(batch_size):
            # Step 2 trains the detector on the step-1 RPN's proposals as FIXED input.
            # reg_targets are built from these boxes, so without the detach the regression
            # target itself would be differentiable and smooth_l1_loss would backpropagate
            # into the RPN.
            proposals = batch_proposals[i].detach()
            gt_boxes = batch_gt_boxes[i].to(device)
            gt_labels = batch_gt_labels[i].to(device)

            labels, matched_gt_idx = self.assign_proposal_labels(proposals, gt_boxes, gt_labels, background_label=background_label)
            sampled_idx = self.sample_proposals(labels, num_samples, pos_fraction, background_label)

            sampled_proposals = proposals[sampled_idx]
            sampled_labels = labels[sampled_idx]

            if gt_boxes.shape[0] == 0:
                # Every label is background here, so these targets are never read.
                # Indexing gt_boxes with matched_gt_idx would be out of bounds.
                reg_targets = torch.zeros((sampled_idx.numel(), 4), dtype=sampled_proposals.dtype, device=device)   # unused placeholder gets discarded by selecting only foreground labels
            else:
                matched_gt_boxes = gt_boxes[matched_gt_idx[sampled_idx]]
                reg_targets = self.encode_box_targets(sampled_proposals, matched_gt_boxes)

            batch_sampled_proposals.append(sampled_proposals)
            batch_labels.append(sampled_labels)
            batch_reg_targets.append(reg_targets)

        pooled_batch = roi_pool(feature_maps, batch_sampled_proposals, batch_img_height, batch_img_width)
        batch_cls_logits, batch_box_deltas = detection_head(pooled_batch)

        total_cls_loss = 0.0
        total_reg_loss = 0.0

        for i in range(batch_size):
            cls_logits = batch_cls_logits[i]
            box_deltas = batch_box_deltas[i]
            labels = batch_labels[i]
            reg_targets = batch_reg_targets[i]

            num_sampled = labels.numel()
            if num_sampled == 0:
                # Nothing survived sampling for this image; it contributes zero to both terms
                continue

            cls_loss = nn.functional.cross_entropy(cls_logits, labels, reduction='mean')

            # For reg_loss, selecting sampled proposals that are not background (i.e., positive samples) is necessary because the regression loss is only computed for foreground classes.
            positive_mask = labels != background_label

            if positive_mask.sum() == 0:
                # smooth_l1_loss on an empty tensor returns nan and would poison the batch
                reg_loss = torch.tensor(0.0, dtype=box_deltas.dtype, device=device)
            else:
                # box_deltas is [N, 4 * 20] 
                num_classes = box_deltas.shape[1] // 4
                positive_labels = labels[positive_mask]
                positive_deltas = box_deltas[positive_mask].view(-1, num_classes, 4)
                # Select the predicted deltas for each proposal based on proposal idx and proposal_label(for that prop_idx)
                # torch.arange(positive_labels.numel(), device=device) creates [0, ..., P = num_positive_proposal], positive_labels = [cls_label_prop_1, cls_label_prop_2, ..., cls_label_prop_P]
                predicted_deltas = positive_deltas[torch.arange(positive_labels.numel(), device=device), positive_labels]

                reg_loss = nn.functional.smooth_l1_loss(predicted_deltas, reg_targets[positive_mask], reduction='sum')

                reg_loss = reg_loss / num_sampled

            total_cls_loss += cls_loss
            total_reg_loss += reg_loss

        return total_cls_loss / batch_size, total_reg_loss / batch_size

    def compute_iou_matrix_corners(self, boxes_a, boxes_b):
        a_x1, a_y1, a_x2, a_y2 = boxes_a[:, 0], boxes_a[:, 1], boxes_a[:, 2], boxes_a[:, 3]
        b_x1, b_y1, b_x2, b_y2 = boxes_b[:, 0], boxes_b[:, 1], boxes_b[:, 2], boxes_b[:, 3]

        inter_x1 = torch.max(a_x1.unsqueeze(1), b_x1.unsqueeze(0))
        inter_y1 = torch.max(a_y1.unsqueeze(1), b_y1.unsqueeze(0))
        inter_x2 = torch.min(a_x2.unsqueeze(1), b_x2.unsqueeze(0))
        inter_y2 = torch.min(a_y2.unsqueeze(1), b_y2.unsqueeze(0))

        inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

        area_a = (a_x2 - a_x1) * (a_y2 - a_y1)
        area_b = (b_x2 - b_x1) * (b_y2 - b_y1)

        union_area = area_a.unsqueeze(1) + area_b.unsqueeze(0) - inter_area

        return inter_area / union_area

    def assign_proposal_labels(self, proposals, gt_boxes, gt_labels, pos_iou_thresh=0.5, neg_iou_lo=0.1, background_label=20):
        # proposals: [N, 4] CORNER absolute px, straight from RegionProposalNetwork
        # gt_boxes:  [M, 4] CORNER absolute px    gt_labels: [M] in 0..19
        # Returns labels [N] (0..19 foreground, 20 background, -1 ignore) and
        # matched_gt_idx [N] (index of the best-overlapping gt box for each proposal).
        num_proposals = proposals.shape[0]

        if gt_boxes.shape[0] == 0:
            # No objects in this image, so every proposal is background
            labels = torch.full((num_proposals,), background_label, dtype=torch.long, device=proposals.device)
            matched_gt_idx = torch.zeros((num_proposals,), dtype=torch.long, device=proposals.device)  # unused placeholder
            return labels, matched_gt_idx

        iou_matrix = self.compute_iou_matrix_corners(proposals, gt_boxes)
        max_iou_per_proposal, matched_gt_idx = iou_matrix.max(dim=1)

        labels = torch.full((num_proposals,), -1, dtype=torch.long, device=proposals.device)
        labels[max_iou_per_proposal >= neg_iou_lo] = background_label

        positive_mask = max_iou_per_proposal >= pos_iou_thresh
        labels[positive_mask] = gt_labels[matched_gt_idx[positive_mask]]

        return labels, matched_gt_idx

    def sample_proposals(self, labels, num_samples=64, pos_fraction=0.25, background_label=20):
        # 64 per image, 25% of them positive.
        positive_idx = torch.where((labels >= 0) & (labels != background_label))[0]
        negative_idx = torch.where(labels == background_label)[0]

        num_pos = min(int(num_samples * pos_fraction), positive_idx.numel())
        num_neg = min(num_samples - num_pos, negative_idx.numel())

        perm_pos = torch.randperm(positive_idx.numel(), device=labels.device)[:num_pos]
        perm_neg = torch.randperm(negative_idx.numel(), device=labels.device)[:num_neg]

        sampled_pos_idx = positive_idx[perm_pos]
        sampled_neg_idx = negative_idx[perm_neg]

        # positives first, then negatives -- ignored proposals never enter the result
        return torch.cat([sampled_pos_idx, sampled_neg_idx])

    def corners_to_center(self, boxes):
        # (x1, y1, x2, y2) -> (x_c, y_c, w, h)
        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
        w = x2 - x1
        h = y2 - y1
        x_c = x1 + w / 2
        y_c = y1 + h / 2
        return torch.stack([x_c, y_c, w, h], dim=1)

    def encode_box_targets(self, proposals, gt_boxes, delta_std=(0.1, 0.1, 0.2, 0.2)):
        proposals_centered = self.corners_to_center(proposals)
        gt_centered = self.corners_to_center(gt_boxes)

        p_xc, p_yc, p_w, p_h = proposals_centered[:, 0], proposals_centered[:, 1], proposals_centered[:, 2], proposals_centered[:, 3]
        gt_xc, gt_yc, gt_w, gt_h = gt_centered[:, 0], gt_centered[:, 1], gt_centered[:, 2], gt_centered[:, 3]

        target_dx = (gt_xc - p_xc) / p_w
        target_dy = (gt_yc - p_yc) / p_h
        target_dw = torch.log(gt_w / p_w)
        target_dh = torch.log(gt_h / p_h)

        target_deltas = torch.stack((target_dx, target_dy, target_dw, target_dh), dim=1)

        # Standard Deviations of the regression targets
        std = torch.tensor(delta_std, dtype=target_deltas.dtype, device=target_deltas.device)

        # Based on the Dataset, the mean of target_deltas is appx 0 and std dev (0.1, 0.1, 0.2, 0.2)
        # So, by (target_deltas - 0)/std we are normalizing the regression targets to have zero mean and unit variance, as described in the Fast R-CNN paper.
        return target_deltas / std


class DetectionNet(nn.Module):
    def __init__(self, detection_head, background_label=20, delta_std=(0.1, 0.1, 0.2, 0.2),
                 score_thresh=0.05, nms_iou_thresh=0.5, max_detections_per_image=100,
                 min_box_size=1.0):
        super().__init__()

        self.detection_head = detection_head
        self.background_label = background_label
        self.delta_std = delta_std

        # Inference-time detection settings. score_thresh=0.05 and 100 detections per
        # image are the Fast R-CNN evaluation conventions; both are deliberately low
        # bars, since AP wants the full precision/recall curve rather than one
        # conservative operating point.
        self.score_thresh = score_thresh
        self.nms_iou_thresh = nms_iou_thresh
        self.max_detections_per_image = max_detections_per_image
        self.min_box_size = min_box_size

    def corners_to_center(self, boxes):
        # (x1, y1, x2, y2) -> (x_c, y_c, w, h)
        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
        w = x2 - x1
        h = y2 - y1
        x_c = x1 + w / 2
        y_c = y1 + h / 2
        return torch.stack([x_c, y_c, w, h], dim=1)

    def select_foreground(self, cls_logits):
        # score per proposal per class, excluding the background class (last column)
        softmax_scores = nn.functional.softmax(cls_logits, dim=-1)      # [N, num_classes + 1]
        fg_scores = softmax_scores[:, :self.background_label]           # [N, num_classes] -- background class is not a foreground class

        proposal_idx, labels = (fg_scores > self.score_thresh).nonzero(as_tuple=True)       # proposal_idx: [K], labels: [K] -- K is the number of proposals that survived score_thresh for at least one class
        scores = fg_scores[proposal_idx, labels]

        return proposal_idx, labels, scores

    def decode_box_deltas(self, proposals_center, box_deltas):
        # proposals_center: [N, 4] (x_c, y_c, w, h)
        # box_deltas: [N, 4] (dx, dy, dw, dh), normalized by self.delta_std (Fast R-CNN convention, matches DetectionLoss.encode_box_targets) -- un-normalize before applying the inverse transform
        std = torch.tensor(self.delta_std, dtype=box_deltas.dtype, device=box_deltas.device)
        dx, dy, dw, dh = (box_deltas * std).unbind(dim=1)

        xc_p, yc_p, w_p, h_p = proposals_center[:, 0], proposals_center[:, 1], proposals_center[:, 2], proposals_center[:, 3]

        xc = dx * w_p + xc_p
        yc = dy * h_p + yc_p
        w = torch.exp(dw) * w_p
        h = torch.exp(dh) * h_p

        xmin = xc - w / 2
        ymin = yc - h / 2
        xmax = xc + w / 2
        ymax = yc + h / 2

        decoded_boxes = torch.stack((xmin, ymin, xmax, ymax), dim=1)
        return decoded_boxes

    def clip_boxes_to_image(self, decoded_boxes, img_height, img_width):
        # decoded_boxes: [N, 4] CORNER absolute px -- same clamp as RegionProposalNetwork.clip_boxes_to_image
        xmin = torch.clamp(decoded_boxes[:, 0], min=0, max=img_width - 1)
        ymin = torch.clamp(decoded_boxes[:, 1], min=0, max=img_height - 1)
        xmax = torch.clamp(decoded_boxes[:, 2], min=0, max=img_width - 1)
        ymax = torch.clamp(decoded_boxes[:, 3], min=0, max=img_height - 1)

        clipped_boxes = torch.stack((xmin, ymin, xmax, ymax), dim=1)
        return clipped_boxes

    def forward(self, rpn_proposals, pooled_proposals, img_sizes_before_pad):
        # rpn_proposals:    list of B x [N_i, 4] CORNER absolute px, from RegionProposalNetwork
        # pooled_proposals: list of B x [N_i, 1024, 7, 7], RoIPool output for the same proposals
        batch_cls_logits, batch_box_deltas = self.detection_head(pooled_proposals)

        batch_size = len(rpn_proposals)

        labels_list = []
        scores_list = []
        boxes_list = []

        for i in range(batch_size):
            img_height, img_width = img_sizes_before_pad[i]

            cls_logits = batch_cls_logits[i]
            box_deltas = batch_box_deltas[i]
            proposals = rpn_proposals[i]

            proposal_idx, labels, scores = self.select_foreground(cls_logits)

            if proposal_idx.numel() == 0:
                # Nothing in this image cleared score_thresh for any class
                labels_list.append(labels)
                scores_list.append(scores)
                boxes_list.append(proposals.new_zeros((0, 4)))
                continue

            # box_deltas is [N, 4 * num_classes] -- reshape to [N, num_classes, 4] 
            num_classes = box_deltas.shape[1] // 4
            selected_proposals = proposals[proposal_idx]
            predicted_deltas = box_deltas.view(-1, num_classes, 4)[proposal_idx, labels]

            proposals_center = self.corners_to_center(selected_proposals)
            decoded_boxes = self.decode_box_deltas(proposals_center, predicted_deltas)
            clipped_boxes = self.clip_boxes_to_image(decoded_boxes, img_height, img_width)

            # Remove boxes that are too small to be valid detections
            widths = clipped_boxes[:, 2] - clipped_boxes[:, 0]
            heights = clipped_boxes[:, 3] - clipped_boxes[:, 1]
            keep = (widths >= self.min_box_size) & (heights >= self.min_box_size)
            clipped_boxes, labels, scores = clipped_boxes[keep], labels[keep], scores[keep]


            if labels.numel() == 0:
                labels_list.append(labels)
                scores_list.append(scores)
                boxes_list.append(clipped_boxes)
                continue

            # Per-CLASS NMS -- boxes of different classes must not suppress each other
            keep_idx = ops.batched_nms(clipped_boxes, scores, labels, self.nms_iou_thresh)
            keep_idx = keep_idx[:self.max_detections_per_image]

            labels_list.append(labels[keep_idx])
            scores_list.append(scores[keep_idx])
            boxes_list.append(clipped_boxes[keep_idx])

        return labels_list, scores_list, boxes_list