Handling of defaultdicts

Author: datumbox

test/test_prototype_transforms_consistency.py

@@ -2,6 +2,7 @@
 import inspect
 import random
 import re
+from collections import defaultdict
 from importlib.machinery import SourceFileLoader
 from pathlib import Path
 
@@ -1039,9 +1040,7 @@ def check(self, t, t_ref, data_kwargs=None):
                 seg_transforms.RandomCrop(size=480),
                 prototype_transforms.Compose(
                     [
-                        PadIfSmaller(
-                            size=480, fill={features.Mask: 255, features.Image: 0, PIL.Image.Image: 0, torch.Tensor: 0}
-                        ),
+                        PadIfSmaller(size=480, fill=defaultdict(lambda: 0, {features.Mask: 255})),
                         prototype_transforms.RandomCrop(size=480),
                     ]
                 ),

torchvision/prototype/transforms/_utils.py

@@ -33,12 +33,13 @@ def _setup_float_or_seq(arg: Union[float, Sequence[float]], name: str, req_size:
 
 
 def _check_fill_arg(fill: Union[FillType, Dict[Type, FillType]]) -> None:
-    if type(fill) == dict:
-        # Do exact type check to avoid accepting default dicts from the user. DefaultDict values can be verified only
-        # at runtime not at construction type.
+    if isinstance(fill, dict):
         for key, value in fill.items():
             # Check key for type
             _check_fill_arg(value)
+        if isinstance(fill, defaultdict) and callable(fill.default_factory):
+            default_value = fill.default_factory()
+            _check_fill_arg(default_value)
     else:
         if fill is not None and not isinstance(fill, (numbers.Number, tuple, list)):
             raise TypeError("Got inappropriate fill arg, only Numbers, tuples, lists and dicts are allowed.")
@@ -75,10 +76,13 @@ def _setup_fill_arg(fill: Union[FillType, Dict[Type, FillType]]) -> Dict[Type, F
     _check_fill_arg(fill)
 
     if isinstance(fill, dict):
-        fill_copy = {}
         for k, v in fill.items():
-            fill_copy[k] = _convert_fill_arg(v)
-        return fill_copy
+            fill[k] = _convert_fill_arg(v)
+        if isinstance(fill, defaultdict) and callable(fill.default_factory):
+            default_value = fill.default_factory()
+            sanitized_default = _convert_fill_arg(default_value)
+            fill.default_factory = functools.partial(_default_arg, sanitized_default)
+        return fill  # type: ignore[return-value]
 
     return _get_defaultdict(_convert_fill_arg(fill))