def test_nms(self, model_type, tmp_path): torch.random.manual_seed(42) image_size = 512 num_classes = 1 batch_size = 2 strides = [8, 16, 32, 64, 128] sizes = [(image_size // stride, ) * 2 for stride in strides] img = torch.zeros(batch_size, 1, image_size, image_size) pred_cls = [ torch.rand(batch_size, num_classes, *size).sub_(0.45).clamp_min_(0) for size in sizes ] pred_cls[0].clamp_min(0.111) pred_reg = [ torch.rand(batch_size, 4, *size).mul_(image_size / 4).round_().clamp_min_(24) for size in sizes ] pred_centerness = [torch.rand(batch_size, 1, *size) for size in sizes] final_pred = model_type.postprocess(pred_cls, pred_reg, pred_centerness, strides, threshold=0.5, nms_threshold=None) final_boxes = final_pred[..., :4] final_scores = final_pred[..., 4:5] final_cls = final_pred[..., 5:] img_with_box = visualize_bbox(img, final_boxes, final_scores, final_cls) subpath = Path( self.DEST, "fcos_targets") if self.DEST is not None else Path(tmp_path) if subpath.is_dir(): subpath.mkdir(exist_ok=True) for i, item in enumerate(img_with_box): filename = os.path.join(subpath, f"created_targets_no_nms_{i}.png") self.save(filename, item) final_pred = model_type.postprocess(pred_cls, pred_reg, pred_centerness, strides, threshold=0.5, nms_threshold=0.01) final_boxes = final_pred[..., :4] final_scores = final_pred[..., 4:5] final_cls = final_pred[..., 5:] img_with_box = visualize_bbox(img, final_boxes, final_cls, final_scores) for i, item in enumerate(img_with_box): filename = os.path.join(subpath, f"created_targets_nms_{i}.png") self.save(filename, item)
def test_class_names(self, img, label, bbox): if label is None: pytest.skip() class_names = {1: "foo", 2: "bar"} no_names = visualize_bbox(img, bbox, label) names = visualize_bbox(img, bbox, label, class_names=class_names) assert names.shape == no_names.shape and not torch.allclose( torch.as_tensor(names), torch.as_tensor(no_names)) if self.DEST is not None and img.ndim == 3: dest = os.path.join(self.DEST, "test_class_names.png") self.save(dest, names)
def test_multiple_scores(self, img, label, bbox, scores): torch.random.manual_seed(42) if scores is None: pytest.skip() tensor1 = torch.rand_like(torch.as_tensor(scores)) tensor2 = torch.rand_like(torch.as_tensor(scores)) tensor = torch.cat([tensor1, tensor2], dim=-1) scores1 = visualize_bbox(img, bbox, label, scores=tensor1) scores2 = visualize_bbox(img, bbox, label, scores=tensor) assert scores1.shape == scores2.shape and not torch.allclose( torch.as_tensor(scores1), torch.as_tensor(scores2)) if self.DEST is not None and img.ndim == 3: dest = os.path.join(self.DEST, "test_multiple_scores.png") self.save(dest, scores2)
def bbox_overlay(img: Tensor, keypoint_dict: Dict[str, Tensor], **kwargs) -> Tensor: coords = keypoint_dict["coords"] cls = keypoint_dict.get("class", None) score = keypoint_dict.get("score", None) result = visualize_bbox(img, coords, cls, score, thickness=1, **kwargs) return result
def test_save_output(self, center_radius, tmp_path): image_size = 512 num_classes = 2 target_bbox = torch.tensor( [ [140, 140, 144, 144], [10, 10, 128, 128], [32, 64, 128, 256], [250, 10, 250 + 31, 10 + 19], [256, 256, 400, 512], ] ) img = torch.zeros(1, image_size, image_size) target_cls = torch.tensor([1, 0, 1, 1, 0]).unsqueeze_(-1) strides = (8, 16, 32, 64, 128) sizes: Tuple[Tuple[int, int]] = tuple((image_size // stride,) * 2 for stride in strides) # type: ignore criterion = FCOSLoss(strides, num_classes, radius=center_radius) targets = criterion.create_targets(target_bbox, target_cls, sizes) cls_targets = [t[0] for t in targets] reg_targets = [t[1] for t in targets] centerness_targets = [t[2] for t in targets] reg_targets = [torch.linalg.norm(x.float().clamp_min(0), dim=-3, keepdim=True) for x in reg_targets] reg_targets = [x.div(x.amax(dim=(-1, -2, -3), keepdim=True).clamp_min_(1)) for x in reg_targets] centerness_targets = [x.clamp_min_(0) for x in centerness_targets] img_with_box = visualize_bbox(img, target_bbox, target_cls)[None] subpath = Path(self.DEST, "fcos_targets") if self.DEST is not None else Path(tmp_path) subpath.mkdir(exist_ok=True) subpath = Path(subpath, f"radius_{center_radius}") subpath.mkdir(exist_ok=True) for level in range(len(strides)): image_path = os.path.join(subpath) c = cls_targets[level][None] r = reg_targets[level][None] cent = centerness_targets[level][None] filename = os.path.join(image_path, f"reg_level_{level}.png") self.blend_and_save(filename, r, img_with_box) filename = os.path.join(image_path, f"centerness_level_{level}.png") self.blend_and_save(filename, cent, img_with_box) for cls_idx in range(c.shape[1]): filename = os.path.join(image_path, f"cls_{cls_idx}_level_{level}.png") self.blend_and_save(filename, c[..., cls_idx, :, :][None], img_with_box)
def test_visualize_bbox(self, img, label, bbox, class_names, scores): if not isinstance(img, torch.Tensor): pytest.skip() result = visualize_bbox(img, bbox, label, scores, class_names) assert isinstance(result, torch.Tensor) assert result.shape[-2:] == img.shape[-2:] assert result.shape[-3] == 3 if img.ndim != 2: assert result.ndim == img.ndim if self.DEST is not None and img.ndim == 3: dest = os.path.join(self.DEST, "test_visualize_bbox.png") self.save(dest, result)
def test_save_output(self, model_type, tmp_path): torch.random.manual_seed(42) image_size = 512 num_classes = 2 batch_size = 3 center_radius = 2 target_bbox = (torch.tensor([ [10, 10, 128, 128], [12, 12, 130, 130], [32, 64, 128, 256], [256, 256, 400, 512], ]).unsqueeze_(0).repeat(batch_size, 1, 1)) img = torch.zeros(4, 1, image_size, image_size) target_cls = torch.tensor([0, 0, 1, 1]).unsqueeze_(-1).repeat(batch_size, 1, 1) strides = (8, 16, 32, 64, 128) sizes: Tuple[Tuple[int, int]] = tuple( (image_size // stride, ) * 2 for stride in strides) # type: ignore criterion = FCOSLoss(strides, num_classes, radius=center_radius) cls_targets, reg_targets, centerness_targets = criterion.create_targets( target_bbox, target_cls, sizes) final_pred = model_type.postprocess(cls_targets, reg_targets, centerness_targets, strides, threshold=0.5) final_boxes = final_pred[..., :4] final_scores = final_pred[..., 4:5] final_cls = final_pred[..., 5:] img_with_box = visualize_bbox(img, final_boxes, final_scores, final_cls) subpath = Path( self.DEST, "fcos_targets") if self.DEST is not None else Path(tmp_path) if subpath.is_dir(): subpath.mkdir(exist_ok=True) for i, item in enumerate(img_with_box): filename = os.path.join(subpath, f"created_targets_{i}.png") self.save(filename, item)
def test_cuda(self, img, label, bbox, class_names, scores): if not isinstance(img, torch.Tensor): pytest.skip() img = img.cuda() bbox = bbox.cuda() label = label.cuda() if label is not None else None scores = scores.cuda() if scores is not None else None result = visualize_bbox(img, bbox, label, scores, class_names) assert isinstance(result, torch.Tensor) assert result.dtype == torch.uint8 assert result.shape[-2:] == img.shape[-2:] assert result.shape[-3] == 3 if img.ndim != 2: assert result.ndim == img.ndim if self.DEST is not None and img.ndim == 3: dest = os.path.join(self.DEST, "test_visualize_bbox.png") self.save(dest, result)
def test_save_output(self, model_type): torch.random.manual_seed(42) image_size = 512 num_classes = 2 batch_size = 3 center_radius = 2 target_bbox = (torch.tensor([ [10, 10, 128, 128], [12, 12, 130, 130], [32, 64, 128, 256], [256, 256, 400, 512], ]).unsqueeze_(0).repeat(batch_size, 1, 1)) img = torch.zeros(4, 1, image_size, image_size) target_cls = torch.tensor([0, 0, 1, 1]).unsqueeze_(-1).repeat(batch_size, 1, 1) strides = [8, 16, 32, 64, 128] sizes = [(image_size // stride, ) * 2 for stride in strides] criterion = FCOSLoss(strides, num_classes, radius=center_radius) cls_targets, reg_targets, centerness_targets = criterion.create_targets( target_bbox, target_cls, sizes) final_pred, _ = model_type.create_boxes(cls_targets, reg_targets, centerness_targets, strides, threshold=0.5) final_boxes = final_pred[..., :4] final_scores = final_pred[..., 4:5] final_cls = final_pred[..., 5:] img_with_box = visualize_bbox(img, final_boxes, final_scores, final_cls) subpath = os.path.join(self.DEST, "fcos_targets") if not os.path.exists(subpath): os.makedirs(subpath) for i, item in enumerate(img_with_box): filename = os.path.join(subpath, f"created_targets_{i}.png") self.save(filename, item)
def bbox_overlay(img: Tensor, keypoint_dict: Dict[str, Any], **kwargs) -> Tensor: coords = keypoint_dict["coords"] cls = keypoint_dict.get("class", None) score = keypoint_dict.get("score", None) names = keypoint_dict.get("names", None) assert coords is None or isinstance(coords, Tensor) assert cls is None or isinstance(cls, Tensor) assert score is None or isinstance(score, Tensor) assert names is None or isinstance(names, Dict) result = visualize_bbox(img, coords, cls, score, names, thickness=1, **kwargs) return result
def test_inputs_unchanged(self, img, label, bbox, class_names, scores): def copy(x): return torch.as_tensor(x).clone() img_c = copy(img) bbox_c = copy(bbox) if label is not None: label_c = copy(label) if scores is not None: scores_c = copy(scores) result = visualize_bbox(img, bbox=bbox, classes=label, scores=scores, class_names=class_names) assert torch.allclose(torch.as_tensor(img), img_c) assert torch.allclose(torch.as_tensor(bbox), bbox_c) if label is not None: assert torch.allclose(torch.as_tensor(label), label_c) if scores is not None: assert torch.allclose(torch.as_tensor(scores), scores_c)