| 24 | |
| 25 | |
| 26 | class CenterNetDecoder(object): |
| 27 | |
| 28 | @staticmethod |
| 29 | def decode_gt(results, batched_inputs): |
| 30 | detections = {} |
| 31 | id_feature = results["id"] if "id" in results else None |
| 32 | fmap = results["fmap"] |
| 33 | device = fmap.device |
| 34 | fmap_h, fmap_w = fmap.shape[2:] |
| 35 | tlbr = (batched_inputs[0]['instances'].gt_boxes.tensor).to(device) / 4 |
| 36 | num_bbox = len(tlbr) |
| 37 | |
| 38 | ## compute center xy cooridnates |
| 39 | cx = (tlbr[:, 0] + tlbr[:, 2]) / 2 |
| 40 | cy = (tlbr[:, 1] + tlbr[:, 3]) / 2 |
| 41 | cx_idx, cy_idx = cx.to(torch.long).clamp(min=0, max=fmap_w-1), cy.to(torch.long).clamp(min=0, max=fmap_h-1) |
| 42 | |
| 43 | ## extract appearance features |
| 44 | if id_feature is not None: |
| 45 | id_feature = F.normalize(id_feature[:, :, cy_idx, cx_idx].permute(0, 2, 1), dim=2) |
| 46 | detections['id_feat'] = id_feature |
| 47 | if fmap is not None: |
| 48 | fmap = F.normalize(fmap[:, :, cy_idx, cx_idx].permute(0, 2, 1), dim=2) |
| 49 | detections['fmap_feat'] = fmap |
| 50 | |
| 51 | detections['bboxes'] = tlbr.reshape(1, num_bbox, 4) |
| 52 | detections['scores'] = torch.ones(1, num_bbox, 1, device=device) |
| 53 | detections['classes'] = torch.zeros(1, num_bbox, 1, device=device, dtype=torch.long) |
| 54 | |
| 55 | return detections |
| 56 | |
| 57 | @staticmethod |
| 58 | def decode(results, cat_spec_wh=False, K=100, tlbr_flag=False, nms_flag=True, return_index=False, |
| 59 | return_box_dict=False, whwh_flag=True, l2norm_flag=True): |
| 60 | r""" |
| 61 | decode output feature map to detection results |
| 62 | |
| 63 | Args: |
| 64 | fmap(Tensor): output feature map |
| 65 | hm(Tensor): output heatmap |
| 66 | wh(Tensor): tensor that represents predicted width-height |
| 67 | reg(Tensor): tensor that represens regression of center points |
| 68 | cat_spec_wh(bool): whether apply gather on tensor `wh` or not |
| 69 | K(int): topk value |
| 70 | """ |
| 71 | hm = results["hm"].sigmoid() |
| 72 | # fmap = results["hm"] |
| 73 | reg = results["reg"] if "reg" in results else None |
| 74 | wh = results["wh"] |
| 75 | id_feature = results["id"] if "id" in results else None |
| 76 | fmap = results["fmap"] if 'fmap' in results else None |
| 77 | detections = {} |
| 78 | |
| 79 | batch, channel, height, width = hm.shape |
| 80 | if nms_flag: |
| 81 | hm = CenterNetDecoder.pseudo_nms(hm) |
| 82 | |
| 83 | scores, index, clses, ys, xs = CenterNetDecoder.topk_score(hm, K=K) |
nothing calls this directly
no outgoing calls
no test coverage detected