Merge branch 'main' into add-vit-swag-huge

Author: YosuaMichael

torchvision/datasets/utils.py

@@ -7,6 +7,7 @@
 import os.path
 import pathlib
 import re
+import sys
 import tarfile
 import urllib
 import urllib.error
@@ -62,7 +63,10 @@ def bar_update(count, block_size, total_size):
 
 
 def calculate_md5(fpath: str, chunk_size: int = 1024 * 1024) -> str:
-    md5 = hashlib.md5()
+    # Setting the `usedforsecurity` flag does not change anything about the functionality, but indicates that we are
+    # not using the MD5 checksum for cryptography. This enables its usage in restricted environments like FIPS. Without
+    # it torchvision.datasets is unusable in these environments since we perform a MD5 check everywhere.
+    md5 = hashlib.md5(**dict(usedforsecurity=False) if sys.version_info >= (3, 9) else dict())
     with open(fpath, "rb") as f:
         for chunk in iter(lambda: f.read(chunk_size), b""):
             md5.update(chunk)

torchvision/prototype/transforms/__init__.py

@@ -21,3 +21,5 @@
 from ._meta import ConvertBoundingBoxFormat, ConvertImageDtype, ConvertImageColorSpace
 from ._misc import Identity, Normalize, ToDtype, Lambda
 from ._type_conversion import DecodeImage, LabelToOneHot
+
+from ._legacy import Grayscale, RandomGrayscale  # usort: skip

torchvision/prototype/transforms/_legacy.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import warnings
+from typing import Any, Dict
+
+from torchvision.prototype.features import ColorSpace
+from torchvision.prototype.transforms import Transform
+from typing_extensions import Literal
+
+from ._meta import ConvertImageColorSpace
+from ._transform import _RandomApplyTransform
+
+
+class Grayscale(Transform):
+    def __init__(self, num_output_channels: Literal[1, 3] = 1) -> None:
+        deprecation_msg = (
+            f"The transform `Grayscale(num_output_channels={num_output_channels})` "
+            f"is deprecated and will be removed in a future release."
+        )
+        if num_output_channels == 1:
+            replacement_msg = (
+                "transforms.ConvertImageColorSpace(old_color_space=ColorSpace.RGB, color_space=ColorSpace.GRAY)"
+            )
+        else:
+            replacement_msg = (
+                "transforms.Compose(\n"
+                "    transforms.ConvertImageColorSpace(old_color_space=ColorSpace.RGB, color_space=ColorSpace.GRAY),\n"
+                "    transforms.ConvertImageColorSpace(old_color_space=ColorSpace.GRAY, color_space=ColorSpace.RGB),\n"
+                ")"
+            )
+        warnings.warn(f"{deprecation_msg} Instead, please use\n\n{replacement_msg}")
+
+        super().__init__()
+        self.num_output_channels = num_output_channels
+        self._rgb_to_gray = ConvertImageColorSpace(old_color_space=ColorSpace.RGB, color_space=ColorSpace.GRAY)
+        self._gray_to_rgb = ConvertImageColorSpace(old_color_space=ColorSpace.GRAY, color_space=ColorSpace.RGB)
+
+    def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
+        output = self._rgb_to_gray(input)
+        if self.num_output_channels == 3:
+            output = self._gray_to_rgb(output)
+        return output
+
+
+class RandomGrayscale(_RandomApplyTransform):
+    def __init__(self, p: float = 0.1) -> None:
+        warnings.warn(
+            "The transform `RandomGrayscale(p=...)` is deprecated and will be removed in a future release. "
+            "Instead, please use\n\n"
+            "transforms.RandomApply(\n"
+            "    transforms.Compose(\n"
+            "        transforms.ConvertImageColorSpace(old_color_space=ColorSpace.RGB, color_space=ColorSpace.GRAY),\n"
+            "        transforms.ConvertImageColorSpace(old_color_space=ColorSpace.GRAY, color_space=ColorSpace.RGB),\n"
+            "    )\n"
+            "    p=...,\n"
+            ")"
+        )
+
+        super().__init__(p=p)
+        self._rgb_to_gray = ConvertImageColorSpace(old_color_space=ColorSpace.RGB, color_space=ColorSpace.GRAY)
+        self._gray_to_rgb = ConvertImageColorSpace(old_color_space=ColorSpace.GRAY, color_space=ColorSpace.RGB)
+
+    def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
+        return self._gray_to_rgb(self._rgb_to_gray(input))

torchvision/prototype/transforms/_meta.py

@@ -64,7 +64,7 @@ def _transform(self, input: Any, params: Dict[str, Any]) -> Any:
         elif is_simple_tensor(input):
             if self.old_color_space is None:
                 raise RuntimeError(
-                    f"In order to convert vanilla tensor images, `{type(self).__name__}(...)` "
+                    f"In order to convert simple tensor images, `{type(self).__name__}(...)` "
                     f"needs to be constructed with the `old_color_space=...` parameter."
                 )