Convert torchvision ColorJitter hue to degrees for fn.color_twist

dali/python/nvidia/dali/experimental/torchvision/v2/color.py

@@ -170,6 +170,9 @@ class ColorJitter(Operator):
         How much to jitter hue. hue_factor is chosen uniformly from [-hue, hue] or the given
         [min, max]. Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. To jitter hue,
         the pixel values of the input image has to be non-negative for conversion to HSV space.
+        DALI rotates hue linearly in YIQ, so results are close to but not identical to
+        ``torchvision.transforms.v2.functional.adjust_hue``, which rotates in HSV. The gap
+        grows with the angle, reaching roughly 24 degrees at a 90 degree rotation.
     device : Literal["cpu", "gpu"], optional, default = "cpu"
         Device to use for the color jitter. Can be ``"cpu"`` or ``"gpu"``.
     """
@@ -212,8 +215,11 @@ def _kernel(self, data_input):
         """
         Performs the color jitter using the ``fn.color_twist`` operator.
         """
+        # torchvision expresses hue as a fraction of a full turn (|hue| <= 0.5), while
+        # fn.color_twist takes the hue delta in degrees.
+        hue_degrees = tuple(float(h) * 360.0 for h in self.hue)
         brightness, contrast, saturation, hue = _get_BrightnessContrastSaturationHue(
-            self.brightness, self.contrast, self.saturation, self.hue, fn.random.uniform
+            self.brightness, self.contrast, self.saturation, hue_degrees, fn.random.uniform
         )
 
         data_input = fn.color_twist(

dali/test/python/torchvision/test_tv_color.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 import os
 
+import numpy as np
 from nose2.tools import params, cartesian_params
 from nose_utils import assert_raises
 from PIL import Image
@@ -149,6 +151,107 @@ def test_colorjitter_images(cj_params, device):
         _ = cj(img)
 
 
+# Mirrors the hue-only path of ref_color_twist in dali/test/python/operator_1/test_color_twist.py.
+# It is duplicated rather than imported because that module lives in another test directory with
+# no package init, and nothing in this tree imports across those directories today.
+_rgb2yiq = np.array([[0.299, 0.587, 0.114], [0.596, -0.274, -0.321], [0.211, -0.523, 0.311]])
+_yiq2rgb = np.linalg.inv(_rgb2yiq)
+
+
+def yiq_hue_rotate(img: np.ndarray, degrees: float) -> np.ndarray:
+    """Rotate hue by `degrees` in YIQ, the way fn.color_twist defines it."""
+    angle = math.radians(degrees)
+    s, c = math.sin(angle), math.cos(angle)
+    hmat = np.array([[1, 0, 0], [0, c, s], [0, -s, c]])
+    m = np.matmul(_yiq2rgb, np.matmul(hmat, _rgb2yiq))
+    pixels = img.reshape([-1, img.shape[-1]]).astype(np.float64)
+    pixels = np.matmul(pixels, m.transpose())
+    return np.clip(np.round(pixels.reshape(img.shape)), 0, 255).astype(np.uint8)
+
+
+def median_hue_shift(before: Image.Image, after: Image.Image) -> float:
+    """Median hue rotation from `before` to `after`, in degrees, over the colorful pixels."""
+    hue_before, saturation, _ = before.convert("HSV").split()
+    hue_after, _, _ = after.convert("HSV").split()
+    to_degrees = 360.0 / 256.0
+    hue_before = np.asarray(hue_before, dtype=np.float64) * to_degrees
+    hue_after = np.asarray(hue_after, dtype=np.float64) * to_degrees
+    # hue is meaningless for near-gray pixels
+    colorful = np.asarray(saturation) > 32
+    assert colorful.any(), "image has no colorful pixels to measure hue on"
+    shift = (hue_after - hue_before + 180.0) % 360.0 - 180.0
+    return float(np.median(shift[colorful]))
+
+
+def hue_error(actual: float, expected: float) -> float:
+    """Absolute difference between two hue rotations, taking the shorter way around."""
+    return abs((actual - expected + 180.0) % 360.0 - 180.0)
+
+
+@cartesian_params((0.05, 0.1, -0.1, 0.25, -0.25, 0.5, -0.5), ("cpu", "gpu"))
+def test_colorjitter_hue_rotation(hue, device):
+    # torchvision expresses hue as a fraction of a full turn, fn.color_twist takes degrees.
+    # DALI rotates linearly in YIQ, which drifts from torchvision's HSV shift in proportion
+    # to the angle, so the tolerance scales too. The comparison is circular because hue=+-0.5
+    # lands on +-180, where implementations that agree closely can report opposite signs.
+    requested = abs(hue) * 360.0
+    # Deliberately loose. test_colorjitter_hue_matches_yiq_reference pins the unit conversion
+    # exactly, so this only has to catch a gross unit error. 0.35 was the measured worst case
+    # and left no margin: at hue=-0.1 the drift is 12.66 degrees against a 12.60 bound. Errors
+    # here move in steps of 360/256 degrees, the PIL HSV hue quantum, so the bound needs more
+    # than one step of slack. 0.45 leaves over 3 degrees everywhere and still catches a 1.5x
+    # unit error.
+    tol = max(4.0, 0.45 * requested)
+    cj = Compose([ColorJitter(hue=(hue, hue), device=device)])
+
+    for fn in test_files:
+        img = Image.open(fn).convert("RGB")
+        expected = median_hue_shift(img, transforms.functional.adjust_hue(img, hue))
+        actual = median_hue_shift(img, cj(img))
+        assert hue_error(actual, expected) < tol, (
+            f"hue={hue} rotated by {actual:.2f} degrees, torchvision rotates by "
+            f"{expected:.2f} degrees: {fn}"
+        )
+
+
+@cartesian_params((0.05, 0.1, -0.1, 0.25, -0.25, 0.5, -0.5), ("cpu", "gpu"))
+def test_colorjitter_hue_matches_yiq_reference(hue, device):
+    # This is what pins the unit conversion. fn.color_twist rotates hue by the number of degrees
+    # it is handed, so ColorJitter(hue=h) has to match a rotation of h*360 degrees in YIQ.
+    # Measuring against DALI's own model takes the YIQ-vs-HSV drift out of the comparison, which
+    # is what forced the loose bound in the test above, and leaves a per-pixel check instead.
+    cj = Compose([ColorJitter(hue=(hue, hue), device=device)])
+
+    for fn in test_files:
+        img = Image.open(fn).convert("RGB")
+        expected = yiq_hue_rotate(np.asarray(img), hue * 360.0)
+        actual = np.asarray(cj(img).convert("RGB"))
+        # The bound test_color_twist.py uses for this operator with uint8 output.
+        assert np.allclose(actual, expected, 1 / 512, 1), (
+            f"hue={hue} on {fn}: max abs difference "
+            f"{np.abs(actual.astype(np.int32) - expected.astype(np.int32)).max()}"
+        )
+
+
+@params("cpu", "gpu")
+def test_colorjitter_hue_range_is_converted(device):
+    # hue=(h, h) short-circuits in _get_BrightnessContrastSaturationHue and passes a scalar.
+    # A genuine range takes the fn.random.uniform branch, which is the one that consumes the
+    # converted range, and is what ColorJitter(hue=0.1) expands to. Both endpoints must be
+    # converted, so the sampled rotation has to land inside [36, 72] degrees.
+    lo, hi = 0.1, 0.2
+    cj = Compose([ColorJitter(hue=(lo, hi), device=device)])
+    tol = max(4.0, 0.45 * hi * 360.0)
+
+    for fn in test_files:
+        img = Image.open(fn).convert("RGB")
+        actual = median_hue_shift(img, cj(img))
+        assert lo * 360.0 - tol <= actual <= hi * 360.0 + tol, (
+            f"hue range ({lo}, {hi}) should rotate within "
+            f"[{lo * 360.0:.0f}, {hi * 360.0:.0f}] degrees, measured {actual:.2f}: {fn}"
+        )
+
+
 """
 TODO (https://github.com/NVIDIA/DALI/issues/DALI-4656):
 DALI ColorJitter does not currently work on CHW layout