import torch import torch.nn as nn import torchvision.ops as ops class RPN_Head(nn.Module): def __init__(self, in_channels, mid_channels): super(RPN_Head, self).__init__() self.num_anchors = 9 self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=1) self.conv_cls = nn.Conv2d(mid_channels, self.num_anchors * 2, kernel_size=1, stride=1) self.conv_reg = nn.Conv2d(mid_channels, self.num_anchors * 4, kernel_size=1, stride=1) self.conv1.weight.data.normal_(0, 0.01) self.conv_cls.weight.data.normal_(0, 0.01) self.conv_reg.weight.data.normal_(0, 0.01) self.conv1.bias.data.zero_() self.conv_cls.bias.data.zero_() self.conv_reg.bias.data.zero_() self.scales = torch.tensor([128, 256, 512]) self.ratios = torch.tensor([0.5, 1, 2]) def forward(self, feature_map, batch_img_height=None, batch_img_width=None): self.img_height = batch_img_height self.img_width = batch_img_width x = torch.relu(self.conv1(feature_map)) batch_cls_logits = self.conv_cls(x).permute(0, 2, 3, 1) batch_box_deltas = self.conv_reg(x).permute(0, 2, 3, 1) batch_cls_logits = batch_cls_logits.reshape(batch_cls_logits.shape[0], -1, 2) batch_box_deltas = batch_box_deltas.reshape(batch_box_deltas.shape[0], -1, 4) batch_anchors = self.generate_anchors(feature_map) return batch_cls_logits, batch_box_deltas, batch_anchors def generate_anchors(self, feature_map): # This function should generate anchors based on the feature map size and predefined scales/aspect ratios batch_size, _, height, width = feature_map.shape # Calculating stride for the feature map relative to the input image total_stride_x = self.img_width // width total_stride_y = self.img_height // height # For each position in the feature map, calculate the center of the anchor box # center_x = (x_coord * total_stride_x) + (0.5 * total_stride_x) # center_y = (y_coord * total_stride_y) + (0.5 * total_stride_y) # Create a grid of center positions grid_x = torch.arange(width).float() + 0.5 grid_y = torch.arange(height).float() + 0.5 # Create a meshgrid of x_center, y_center positions for each anchor for each value of width and height in feature map grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing='ij') # (H, W) format # Converting grid positions to the scale from the feature map to the image size grid_x = grid_x * total_stride_x grid_y = grid_y * total_stride_y # Now, for each center position, we need to generate anchors based on the scales and aspect ratios widths = self.scales.view(-1, 1) * torch.sqrt(self.ratios).view(1, -1) heights = self.scales.view(-1, 1) / torch.sqrt(self.ratios).view(1, -1) centres = torch.stack((grid_x, grid_y), dim=-1).unsqueeze(dim=-2) # Shape: [height, width, 1, 2] sizes = torch.stack((widths.flatten(), heights.flatten()), dim=-1)[None, None, : ] # Shape: [1, 1, num_scales * num_ratios, 2] anchors = torch.cat((centres.expand(-1, -1, sizes.shape[2], -1), sizes.expand(centres.shape[0], centres.shape[1], -1, -1)), dim=-1) # Shape: [height, width, num_anchors, 4] anchors = anchors.view(-1, 4) # Flatten to [num_anchors_total, 4] batch_anchors = anchors.repeat(batch_size, 1, 1) # Repeat for each image in the batch batch_anchors = batch_anchors.to(feature_map.device) # Move to the same device as the feature map return batch_anchors class RPN_Loss(nn.Module): def __init__(self): super(RPN_Loss, self).__init__() def forward(self, batch_cls_logits, batch_box_deltas, batch_anchors, batch_gt_boxes, img_sizes_before_pad): batch_size = batch_cls_logits.shape[0] total_cls_loss = 0.0 total_reg_loss = 0.0 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] anchors = batch_anchors[i] gt_boxes = batch_gt_boxes[i] gt_boxes = gt_boxes.to(anchors.device) # Move gt_boxes to the same device as anchors inside_indices = self.anchors_inside_image(anchors, img_height, img_width) anchors = anchors[inside_indices] cls_logits = cls_logits[inside_indices] box_deltas = box_deltas[inside_indices] if gt_boxes.shape[0] == 0: # No objects in this image: every anchor is negative by definition, # there's nothing to regress toward. anchor_labels = torch.full((anchors.shape[0],), -1, dtype=torch.long, device=anchors.device) matched_gt_indices = torch.zeros((anchors.shape[0],), dtype=torch.long, device=anchors.device) #unused placeholder else: anchor_labels, matched_gt_indices = self.anchor_labelling(anchors, gt_boxes) sample_mask, sampled_pos_idx, sampled_neg_idx = self.create_sample_mask_per_img(anchor_labels) cls_loss = self.cls_loss_fn_per_img(cls_logits, anchor_labels, sample_mask) if sampled_pos_idx.numel() == 0: reg_loss = torch.tensor(0.0, device=box_deltas.device, dtype=box_deltas.dtype) else: reg_loss = self.reg_loss_fn_per_img(box_deltas, gt_boxes, anchors, sampled_pos_idx, matched_gt_indices) total_cls_loss += cls_loss total_reg_loss += reg_loss return total_cls_loss / batch_size, total_reg_loss / batch_size # Return average losses over the batch def cls_loss_fn_per_img(self, cls_logits, anchor_labels, sample_mask): selected_cls_logits = cls_logits[sample_mask] selected_anchor_labels = anchor_labels[sample_mask] targets = (selected_anchor_labels == 1).long() # Convert to binary targets (1 for positive, 0 for negative) cls_loss = nn.functional.cross_entropy(selected_cls_logits, targets, reduction='mean') # reduction='mean' to get the average loss over the selected samples i.e., divide by num_samples = Ncls return cls_loss def corners_to_center(self, boxes): 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 reg_loss_fn_per_img(self, box_deltas, gt_boxes, valid_anchors, sampled_pos_idx, matched_gt_indices, N_reg=2400): pred_box_deltas = box_deltas[sampled_pos_idx] selected_gt_boxes = gt_boxes[matched_gt_indices[sampled_pos_idx]] # Select the corresponding gt_boxes for the sampled positive anchors anchors = valid_anchors[sampled_pos_idx] gt_boxes_centered = self.corners_to_center(selected_gt_boxes) xc_a, yc_a, w_a, h_a = anchors[:, 0], anchors[:, 1], anchors[:, 2], anchors[:, 3] gt_xc, gt_yc, gt_w, gt_h = gt_boxes_centered[:, 0], gt_boxes_centered[:, 1], gt_boxes_centered[:, 2], gt_boxes_centered[:, 3] # Calculate the regression targets (deltas) for the selected anchors target_dx = (gt_xc - xc_a) / w_a target_dy = (gt_yc - yc_a) / h_a target_dw = torch.log(gt_w / w_a) target_dh = torch.log(gt_h / h_a) target_box_deltas = torch.stack((target_dx, target_dy, target_dw, target_dh), dim=1) reg_loss = nn.functional.smooth_l1_loss(pred_box_deltas, target_box_deltas, reduction='sum') reg_loss /= N_reg return reg_loss def create_sample_mask_per_img(self, anchor_labels, num_samples=256, pos_fraction=0.5): positive_idx = torch.where(anchor_labels == 1)[0] # ...[0] as it is a single valued tuple negative_idx = torch.where(anchor_labels == -1)[0] num_pos = min(int(num_samples * pos_fraction), positive_idx.numel()) # numel() returns the number of elements in the tensor num_neg = min(num_samples - num_pos, negative_idx.numel()) perm_pos = torch.randperm(positive_idx.numel(), device=anchor_labels.device)[:num_pos] perm_neg = torch.randperm(negative_idx.numel(), device=anchor_labels.device)[:num_neg] sampled_pos_idx = positive_idx[perm_pos] sampled_neg_idx = negative_idx[perm_neg] # Ignoring i.e., making the anchor labels of the unsampled anchors to 0 (ignore) sample_mask = torch.zeros_like(anchor_labels, dtype=torch.bool, device=anchor_labels.device) sample_mask[sampled_pos_idx] = True sample_mask[sampled_neg_idx] = True return sample_mask, sampled_pos_idx, sampled_neg_idx def anchor_labelling(self, anchors, gt_boxes, pos_iou_threshold=0.7, neg_iou_threshold=0.3): iou_matrix = self.compute_iou_matrix(anchors, gt_boxes) num_anchors, num_gt_boxes = iou_matrix.shape labels = torch.zeros((num_anchors,), dtype=torch.long, device=anchors.device) # Initialize all labels to 0 (ignore) max_iou_per_anchor, matched_gt_indices = iou_matrix.max(dim=1) labels[max_iou_per_anchor < neg_iou_threshold] = -1 # Negative labels labels[max_iou_per_anchor >= pos_iou_threshold] = 1 # Positive labels max_iou_per_gt, matched_anchor_indices = iou_matrix.max(dim=0) labels[matched_anchor_indices] = 1 # Ensure each gt box has at least one positive anchor return labels, matched_gt_indices def anchors_inside_image(self, anchors, img_height, img_width): # Check if anchors are inside the image boundaries xc, yc, w, h = anchors[:, 0], anchors[:, 1], anchors[:, 2], anchors[:, 3] x1 = xc - w / 2 y1 = yc - h / 2 x2 = xc + w / 2 y2 = yc + h / 2 inside_indices = (0<=x1) & (0<=y1) & (x2 <= img_width-1) & (y2 <= img_height-1) # anchors[:, 0] is xc, anchors[:, 1] is yc, anchors[:, 2] is w, anchors[:, 3] is h return inside_indices def compute_iou_matrix(self, anchors, gt_boxes): # Compute the IoU matrix between anchors and ground truth boxes # anchors: [num_anchors, 4], gt_boxes: [num_gt_boxes, 4] # Returns: IoU matrix of shape [num_anchors, num_gt_boxes] # Convert anchors and gt_boxes to (x1, y1, x2, y2) format anchor_x1 = anchors[:, 0] - anchors[:, 2] / 2 anchor_y1 = anchors[:, 1] - anchors[:, 3] / 2 anchor_x2 = anchors[:, 0] + anchors[:, 2] / 2 anchor_y2 = anchors[:, 1] + anchors[:, 3] / 2 gt_x1 = gt_boxes[:, 0] gt_y1 = gt_boxes[:, 1] gt_x2 = gt_boxes[:, 2] gt_y2 = gt_boxes[:, 3] # Calculate intersection inter_x1 = torch.max(anchor_x1.unsqueeze(1), gt_x1.unsqueeze(0)) # Unsqueeze to align dimensions for broadcasting inter_y1 = torch.max(anchor_y1.unsqueeze(1), gt_y1.unsqueeze(0)) inter_x2 = torch.min(anchor_x2.unsqueeze(1), gt_x2.unsqueeze(0)) inter_y2 = torch.min(anchor_y2.unsqueeze(1), gt_y2.unsqueeze(0)) inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0) # Calculate union anchor_area = (anchor_x2 - anchor_x1) * (anchor_y2 - anchor_y1) gt_area = (gt_x2 - gt_x1) * (gt_y2 - gt_y1) union_area = anchor_area.unsqueeze(dim=1) + gt_area.unsqueeze(dim=0) - inter_area # Unsqueeze to align dimensions for broadcasting # Compute IoU iou_matrix = inter_area / union_area return iou_matrix import torchvision.ops as ops class RegionProposalNetwork(nn.Module): def __init__(self, rpn_head=None, in_channels=1024, mid_channels=512): super(RegionProposalNetwork, self).__init__() if rpn_head is not None: self.rpn_head = rpn_head else: self.rpn_head = RPN_Head(in_channels, mid_channels) def forward(self, feature_map, batch_img_height=None, batch_img_width=None, img_sizes_before_pad=None, pre_nms_top_n=6000, post_nms_top_n=300): # pre_nms_top_n: proposals kept per image (by score) before NMS. # post_nms_top_n: proposals kept per image after NMS. Paper uses 2000 for # Step-2 Fast R-CNN training and 300 for inference/eval; defaults to 300. batch_cls_logits, batch_box_deltas, batch_anchors = self.rpn_head(feature_map, batch_img_height=batch_img_height, batch_img_width=batch_img_width) batch_size = batch_cls_logits.shape[0] 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] anchors = batch_anchors[i] scores = nn.functional.softmax(cls_logits, dim=-1)[:, 1] # Get the foreground class scores for each box decoded_boxes = self.decode_box_deltas(anchors, box_deltas) clipped_boxes = self.clip_boxes_to_image(decoded_boxes, img_height, img_width) filtered_boxes, filtered_scores = self.filter_small_boxes(scores, clipped_boxes, min_size=16) if filtered_boxes.shape[0] == 0: # No proposals survived filtering for this image — return empty, well-shaped tensors scores_list.append(torch.zeros((0,), device=feature_map.device)) # shape (0,) to not break further tensor operations in the pipeline boxes_list.append(torch.zeros((0, 4), device=feature_map.device)) # shape (0, 4) to not break further tensor operations in the pipeline continue pre_nms_top_n_boxes, pre_nms_top_n_scores = self.pre_nms_top_n(filtered_boxes, filtered_scores, pre_nms_top_n=pre_nms_top_n) post_nms_indices = self.nms(pre_nms_top_n_boxes, pre_nms_top_n_scores, iou_threshold=0.7) post_nms_top_n_boxes, post_nms_top_n_scores = self.post_nms_top_n(pre_nms_top_n_boxes, pre_nms_top_n_scores, post_nms_indices, post_nms_top_n=post_nms_top_n) scores_list.append(post_nms_top_n_scores) boxes_list.append(post_nms_top_n_boxes) return scores_list, boxes_list # Decoding box deltas to get the final bounding boxes per img def decode_box_deltas(self, anchors, box_deltas): # Decode the predicted box deltas to get the final bounding boxes # anchors: [num_anchors, 4], box_deltas: [num_anchors, 4] # Returns: decoded_boxes of shape [num_anchors, 4] xc_a, yc_a, w_a, h_a = anchors[:, 0], anchors[:, 1], anchors[:, 2], anchors[:, 3] dx, dy, dw, dh = box_deltas[:, 0], box_deltas[:, 1], box_deltas[:, 2], box_deltas[:, 3] # Apply the inverse of the encoding transformation xc = dx * w_a + xc_a yc = dy * h_a + yc_a w = torch.exp(dw) * w_a h = torch.exp(dh) * h_a # Convert back to (xmin, ymin, xmax, ymax) format 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 # Clip the decoded boxes to ensure they are within the image boundaries per img def clip_boxes_to_image(self, decoded_boxes, img_height, img_width): # Clip the decoded boxes to ensure they are within the image boundaries # decoded_boxes: [num_anchors, 4] # Returns: clipped_boxes of shape [num_anchors, 4] 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 filter_small_boxes(self, scores, decoded_boxes, min_size=16): # Filter out boxes that are smaller than a specified minimum size # decoded_boxes: [num_anchors, 4] # Returns: filtered_boxes of shape [num_filtered_anchors, 4] widths = decoded_boxes[:, 2] - decoded_boxes[:, 0] heights = decoded_boxes[:, 3] - decoded_boxes[:, 1] keep_indices = (widths >= min_size) & (heights >= min_size) filtered_boxes = decoded_boxes[keep_indices] # also need to remove the scores corresponding to the filtered boxes, but that will be handled in the pre_nms_top_n function filtered_scores = scores[keep_indices] # Assuming scores is available in the context where this function is called return filtered_boxes, filtered_scores # Get the top N boxes based on their scores before applying NMS per img def pre_nms_top_n(self, decoded_boxes, scores, pre_nms_top_n=6000): # Select the top N boxes based on their scores before applying NMS # decoded_boxes: [num_anchors, 4], scores: [num_anchors] # Returns: top_decoded_boxes of shape [pre_nms_top_n, 4], top_scores of shape [pre_nms_top_n] # scores = nn.functional.softmax(scores, dim=-1)[:1] # Get the foreground class scores for each box, since 1:obj, 0: background, therefore [:1] gives the foreground class scores for each box top_scores, top_indices = torch.topk(scores, k=min(pre_nms_top_n, scores.size(0))) top_decoded_boxes = decoded_boxes[top_indices] return top_decoded_boxes, top_scores def nms(self, decoded_boxes, scores, iou_threshold=0.7): # Apply Non-Maximum Suppression (NMS) to remove overlapping boxes # decoded_boxes: [num_anchors, 4], scores: [num_anchors] # Returns: keep_indices of shape [num_kept_anchors] keep_indices = ops.nms(decoded_boxes, scores, iou_threshold) return keep_indices def post_nms_top_n(self, bboxes, scores, keep_indices, post_nms_top_n=2000): # Select the top N boxes after applying NMS # decoded_boxes_post_nms: [num_kept_anchors, 4], scores_post_nms: [num_kept_anchors] # Returns: final_decoded_boxes of shape [post_nms_top_n, 4], final_scores of shape [post_nms_top_n] scores_post_nms = scores[keep_indices] boxes_post_nms = bboxes[keep_indices] top_scores_post_nms, top_indices_post_nms = torch.topk(scores_post_nms, k=min(post_nms_top_n, scores_post_nms.size(0))) final_decoded_boxes = boxes_post_nms[top_indices_post_nms] return final_decoded_boxes, top_scores_post_nms