[summary] Args: samples (NestedTensor): batch of imgs. B,3,H,W known_boxes (list of knownBox): [knownbox_each_img x B] Returns: [Tensor]: Masked imgs. B,3,H,W.
(samples: NestedTensor, known_boxes)
| 103 | |
| 104 | @torch.no_grad() |
| 105 | def mask_sample(samples: NestedTensor, known_boxes): |
| 106 | """[summary] |
| 107 | |
| 108 | Args: |
| 109 | samples (NestedTensor): batch of imgs. B,3,H,W |
| 110 | known_boxes (list of knownBox): [knownbox_each_img x B] |
| 111 | |
| 112 | Returns: |
| 113 | [Tensor]: Masked imgs. B,3,H,W. |
| 114 | """ |
| 115 | # print("HERE!!!!!!!!!") |
| 116 | # import pdb; pdb.set_trace() |
| 117 | boxes_flat = [ |
| 118 | box_ops.box_cxcywh_to_xyxy(kbs[:, :4]) |
| 119 | for idx, kbs in enumerate(known_boxes) |
| 120 | ] |
| 121 | img_shapes = samples.imgsize() |
| 122 | device = samples.tensors.device |
| 123 | # ! TODO: |
| 124 | for idx, (shape, boxes) in enumerate(zip(img_shapes, boxes_flat)): |
| 125 | h, w = shape.tolist() |
| 126 | scale = torch.Tensor([w, h, w, h]).to(device) |
| 127 | boxes = boxes * scale |
| 128 | for box in boxes: |
| 129 | x1, y1, x2, y2 = [int(i) for i in box.tolist()] |
| 130 | samples.tensors[idx, :, y1:y2, x1:x2] = 0 |
| 131 | return samples |
| 132 | |
| 133 | |
| 134 | class AssignResult(NiceRepr): |