Skip to content

Commit

Permalink
add sequence fill support for ElasticTransform (#7141)
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier authored Jan 27, 2023
1 parent 2bc8a14 commit 71073cb
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 4 deletions.
32 changes: 32 additions & 0 deletions test/test_transforms_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -858,3 +858,35 @@ def test_gaussian_blur(device, channels, meth_kwargs):
agg_method="max",
tol=tol,
)


@pytest.mark.parametrize("device", cpu_and_gpu())
@pytest.mark.parametrize(
"fill",
[
1,
1.0,
[1],
[1.0],
(1,),
(1.0,),
[1, 2, 3],
[1.0, 2.0, 3.0],
(1, 2, 3),
(1.0, 2.0, 3.0),
],
)
@pytest.mark.parametrize("channels", [1, 3])
def test_elastic_transform(device, channels, fill):
if isinstance(fill, (list, tuple)) and len(fill) > 1 and channels == 1:
# For this the test would correctly fail, since the number of channels in the image does not match `fill`.
# Thus, this is not an issue in the transform, but rather a problem of parametrization that just gives the
# product of `fill` and `channels`.
return

_test_class_op(
T.ElasticTransform,
meth_kwargs=dict(fill=fill),
channels=channels,
device=device,
)
2 changes: 0 additions & 2 deletions torchvision/transforms/functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -1539,8 +1539,6 @@ def elastic_transform(
fill (number or str or tuple): Pixel fill value for constant fill. Default is 0.
If a tuple of length 3, it is used to fill R, G, B channels respectively.
This value is only used when the padding_mode is constant.
Only number is supported for torch Tensor.
Only int or str or tuple value is supported for PIL Image.
"""
if not torch.jit.is_scripting() and not torch.jit.is_tracing():
_log_api_usage_once(elastic_transform)
Expand Down
8 changes: 6 additions & 2 deletions torchvision/transforms/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -2104,8 +2104,12 @@ def __init__(self, alpha=50.0, sigma=5.0, interpolation=InterpolationMode.BILINE
interpolation = _interpolation_modes_from_int(interpolation)
self.interpolation = interpolation

if not isinstance(fill, (int, float)):
raise TypeError(f"fill should be int or float. Got {type(fill)}")
if isinstance(fill, (int, float)):
fill = [float(fill)]
elif isinstance(fill, (list, tuple)):
fill = [float(f) for f in fill]
else:
raise TypeError(f"fill should be int or float or a list or tuple of them. Got {type(fill)}")
self.fill = fill

@staticmethod
Expand Down

0 comments on commit 71073cb

Please sign in to comment.