ROI Align from image features using box coordinates in CxCyWH [0,1] format. Args: img_feats_nchw: [B, C, H, W] backbone feature map boxes_cxcywh: [N_boxes, B, 4] in normalized CxCyWH format roi_size: output spatial size Returns: roi_pooled: [N_boxes*B,
(img_feats_nchw, boxes_cxcywh, roi_size=7)
| 126 | # ── ROI Align implementation ──────────────────────────────────────────────── |
| 127 | |
| 128 | def do_roi_align(img_feats_nchw, boxes_cxcywh, roi_size=7): |
| 129 | """ |
| 130 | ROI Align from image features using box coordinates in CxCyWH [0,1] format. |
| 131 | |
| 132 | Args: |
| 133 | img_feats_nchw: [B, C, H, W] backbone feature map |
| 134 | boxes_cxcywh: [N_boxes, B, 4] in normalized CxCyWH format |
| 135 | roi_size: output spatial size |
| 136 | |
| 137 | Returns: |
| 138 | roi_pooled: [N_boxes*B, C, roi_size, roi_size] |
| 139 | """ |
| 140 | n_boxes, bs, _ = boxes_cxcywh.shape |
| 141 | H, W = img_feats_nchw.shape[-2:] |
| 142 | |
| 143 | boxes_xyxy = box_cxcywh_to_xyxy(boxes_cxcywh) |
| 144 | scale = torch.tensor([W, H, W, H], dtype=boxes_xyxy.dtype, device=boxes_xyxy.device) |
| 145 | scale = scale.view(1, 1, 4) |
| 146 | boxes_xyxy = boxes_xyxy * scale |
| 147 | |
| 148 | # roi_align expects list of [N_boxes, 4] per batch element |
| 149 | rois_list = boxes_xyxy.float().transpose(0, 1).unbind(0) |
| 150 | sampled = torchvision.ops.roi_align(img_feats_nchw, rois_list, roi_size) |
| 151 | return sampled # [B*N_boxes, C, roi_size, roi_size] |
| 152 | |
| 153 | |
| 154 | # ── Main dump logic ─────────────────────────────────────────────────────── |
no test coverage detected