diff --git a/visualDet3D/networks/pipelines/testers.py b/visualDet3D/networks/pipelines/testers.py index efa3925..f0e9482 100644 --- a/visualDet3D/networks/pipelines/testers.py +++ b/visualDet3D/networks/pipelines/testers.py @@ -39,4 +39,52 @@ def test_stereo_detection(data, module:nn.Module, scores, bbox, obj_index = module([left_images.cuda().float().contiguous(), right_images.cuda().float().contiguous(), torch.tensor(P2).cuda().float(), torch.tensor(P3).cuda().float()]) obj_types = [cfg.obj_types[i.item()] for i in obj_index] - return scores, bbox, obj_types \ No newline at end of file + return scores, bbox, obj_types + +@PIPELINE_DICT.register_module +@torch.no_grad() +def test_mono_detection_diggingmono3d(data, module:nn.Module, + writer:SummaryWriter, + loss_logger:LossLogger=None, + global_step:int=None, + cfg:EasyDict=EasyDict())-> Tuple[torch.Tensor, torch.Tensor, List[str]]: + image, P2 = data[0], data[1] + + scores, bbox, obj_index = module( + [image.cuda().float().contiguous(), torch.tensor(P2).cuda().float()]) + + if len(scores) == 0: + return scores, bbox, obj_index + + expanded_scores = [] + expanded_bboxes = [] + expanded_obj_index = [] + + for i in range(len(scores)): + expanded_scores.append(scores[i]) + expanded_bboxes.append(bbox[i]) + expanded_obj_index.append(obj_index[i]) + depth = bbox[i, 6] + if depth < 10: + # no augmentation for these measurements + continue + + lamda = 80 + for depth_shift in [-2, -1, -0.5, 0.5, 1, 2]: + new_depth = depth + depth_shift + new_bbox = bbox[i].clone() + new_bbox[6] = new_depth + sigma = torch.exp(depth / lamda) + t_s = torch.exp( - depth_shift**2 / sigma**2) + new_score = scores[i] * t_s + expanded_scores.append(new_score) + expanded_bboxes.append(new_bbox) + expanded_obj_index.append(obj_index[i]) + + expanded_scores = torch.stack(expanded_scores) + expanded_bboxes = torch.stack(expanded_bboxes, dim=0) + + + obj_types = [cfg.obj_types[i.item()] for i in expanded_obj_index] + + return expanded_scores, expanded_bboxes, obj_types