diff --git a/test/test_models.py b/test/test_models.py index 1c0b4892209..24b5a8b6b66 100644 --- a/test/test_models.py +++ b/test/test_models.py @@ -248,6 +248,25 @@ def test_fasterrcnn_switch_devices(self): self.assertTrue("scores" in out_cpu[0]) self.assertTrue("labels" in out_cpu[0]) + def test_generalizedrcnn_transform_repr(self): + + min_size, max_size = 224, 299 + image_mean = [0.485, 0.456, 0.406] + image_std = [0.229, 0.224, 0.225] + + t = models.detection.transform.GeneralizedRCNNTransform(min_size=min_size, + max_size=max_size, + image_mean=image_mean, + image_std=image_std) + + # Check integrity of object __repr__ attribute + expected_string = 'GeneralizedRCNNTransform(' + _indent = '\n ' + expected_string += '{0}Normalize(mean={1}, std={2})'.format(_indent, image_mean, image_std) + expected_string += '{0}Resize(min_size=({1},), max_size={2}, '.format(_indent, min_size, max_size) + expected_string += "mode='bilinear')\n)" + self.assertEqual(t.__repr__(), expected_string) + for model_name in get_available_classification_models(): # for-loop bodies don't define scopes, so we have to save the variables diff --git a/torchvision/models/detection/transform.py b/torchvision/models/detection/transform.py index f1cf8a41bfd..9b2ef009cb8 100644 --- a/torchvision/models/detection/transform.py +++ b/torchvision/models/detection/transform.py @@ -181,6 +181,15 @@ def postprocess(self, result, image_shapes, original_image_sizes): result[i]["keypoints"] = keypoints return result + def __repr__(self): + format_string = self.__class__.__name__ + '(' + _indent = '\n ' + format_string += "{0}Normalize(mean={1}, std={2})".format(_indent, self.image_mean, self.image_std) + format_string += "{0}Resize(min_size={1}, max_size={2}, mode='bilinear')".format(_indent, self.min_size, + self.max_size) + format_string += '\n)' + return format_string + def resize_keypoints(keypoints, original_size, new_size): # type: (Tensor, List[int], List[int])