feat(transforms): add 8 robotics-relevant image augmentations (#4210)

* feat(transforms): add 8 robotics-relevant image augmentations Add GaussianNoise, MotionBlur, JPEGCompression, GaussianPatchBrightness, RandomShadow, CoarseDropout, GammaCorrection, and PlanckianJitter. Each transform addresses a real-world failure mode not covered by the existing 6 defaults (sensor noise, motion blur, compression artifacts, uneven lighting, cast shadows, partial occlusion, exposure variation, color temperature shift). All transforms are pure PyTorch, follow the make_params/transform pattern, and integrate with ImageTransformConfig via a registry. * add augmentation showcase image for PR * update showcase with better sample frame * tune showcase to balanced augmentation intensity * tune showcase: softer shadow, dropout, jitter intensity * refactor(transforms): several updates * update image * chore(media): remove example * chore: add link to example --------- Co-authored-by: Yuxian LI <liyuxian1358@gmail.com>
2026-07-31 21:49:45 +00:00 · 2026-07-31 14:12:45 +02:00
parent af4d15f9ac
commit 414d0eecbd
3 changed files with 647 additions and 3 deletions
@@ -28,9 +28,17 @@ from lerobot.scripts.lerobot_imgtransform_viz import (
    save_each_transform,
 )
 from lerobot.transforms import (
+    CoarseDropout,
+    GammaCorrection,
+    GaussianNoise,
+    GaussianPatchBrightness,
    ImageTransformConfig,
    ImageTransforms,
    ImageTransformsConfig,
+    JPEGCompression,
+    MotionBlur,
+    PlanckianJitter,
+    RandomShadow,
    RandomSubsetApply,
    SharpnessJitter,
    make_transform_from_config,
@@ -455,3 +463,153 @@ def test_save_each_transform(img_tensor_factory, tmp_path):
            assert (transform_dir / file_name).exists(), (
                f"{file_name} was not found in {transform} directory."
            )
+
+
+# --- Tests for robotics-relevant augmentations ---
+
+ROBOTICS_TRANSFORMS = [
+    ("GaussianNoise", GaussianNoise, {"std": (5.0, 25.0)}),
+    ("MotionBlur", MotionBlur, {"kernel_size": (3, 11)}),
+    ("JPEGCompression", JPEGCompression, {"quality": (15, 75)}),
+    ("GaussianPatchBrightness", GaussianPatchBrightness, {}),
+    ("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}),
+    ("CoarseDropout", CoarseDropout, {"max_holes": 8}),
+    ("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}),
+    ("PlanckianJitter", PlanckianJitter, {"temperature": (3_000, 15_000)}),
+]
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_shape_preserved(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_output_range(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.min() >= -0.01, f"{name} min below range: {out.min():.4f}"
+    assert out.max() <= 1.01, f"{name} max above range: {out.max():.4f}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_float_output(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.is_floating_point(), f"{name} output dtype={out.dtype}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_non_float_passthrough(name, cls, kwargs):
+    int_img = torch.randint(0, 255, (3, 32, 32), dtype=torch.uint8)
+    tf = cls(**kwargs)
+    out = tf(int_img)
+    assert torch.equal(out, int_img), f"{name} modified non-float input"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_via_config(name, cls, kwargs):
+    cfg = ImageTransformConfig(type=name, kwargs=kwargs)
+    tf = make_transform_from_config(cfg)
+    assert isinstance(tf, cls), f"Config produced {type(tf)}, expected {cls}"
+
+
+def test_make_transform_error_message_includes_custom():
+    """Error message should list all registered custom transforms."""
+    with pytest.raises(ValueError, match="GaussianNoise"):
+        make_transform_from_config(ImageTransformConfig(type="NonExistent"))
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+@pytest.mark.parametrize("shape", [(4, 3, 32, 32), (2, 4, 3, 16, 16)])
+def test_robotics_transform_supports_temporal_batches(name, cls, kwargs, shape):
+    img = torch.rand(shape)
+    out = cls(**kwargs)(img)
+    assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
+    assert out.min() >= 0
+    assert out.max() <= 1
+
+
+@pytest.mark.parametrize(
+    "cls,kwargs",
+    [
+        (GaussianNoise, {"std": (25.0, 25.0)}),
+        (MotionBlur, {"kernel_size": 5}),
+        (JPEGCompression, {"quality": 10}),
+        (
+            GaussianPatchBrightness,
+            {"num_patches": 1, "sigma_range": (0.2, 0.2), "factor_range": (0.5, 0.5)},
+        ),
+        (RandomShadow, {"opacity": 0.5}),
+        (CoarseDropout, {"max_holes": 1, "fill_value": 0.0}),
+        (GammaCorrection, {"gamma": (2.0, 2.0)}),
+        (PlanckianJitter, {"temperature": 3_000}),
+    ],
+)
+def test_robotics_transform_is_not_silent_noop(cls, kwargs):
+    img = torch.rand(3, 32, 32)
+    out = cls(**kwargs)(img)
+    assert not torch.equal(out, img)
+
+
+@pytest.mark.parametrize(
+    "transform",
+    [
+        GaussianNoise(std=25),
+        RandomShadow(opacity=0.5),
+        CoarseDropout(max_holes=4),
+    ],
+)
+def test_robotics_transform_random_params_are_reused(transform):
+    img = torch.rand(3, 32, 32)
+    params = transform.make_params([img])
+    torch.testing.assert_close(transform.transform(img, params), transform.transform(img, params))
+
+
+def test_motion_blur_kernel_size_stays_in_configured_range():
+    transform = MotionBlur(kernel_size=(4, 10))
+    sampled_sizes = {transform.make_params([])["kernel_size"] for _ in range(100)}
+    assert sampled_sizes <= {5, 7, 9}
+    assert sampled_sizes
+
+
+def test_gamma_correction_scalar_below_one_defines_symmetric_range():
+    transform = GammaCorrection(gamma=0.5)
+    assert transform.gamma == (0.5, 2.0)
+    assert transform(torch.rand(3, 8, 8)).shape == (3, 8, 8)
+
+
+def test_planckian_jitter_uses_correlated_temperature_coefficients():
+    img = torch.full((2, 3, 8, 8), 0.25)
+    out = PlanckianJitter(temperature=3_000)(img)
+    torch.testing.assert_close(out[:, 1], img[:, 1])
+    assert torch.all(out[:, 0] > out[:, 1])
+    assert torch.all(out[:, 2] < out[:, 1])
+
+
+def test_random_shadow_supports_small_images():
+    img = torch.rand(3, 7, 7)
+    assert RandomShadow()(img).shape == img.shape
+
+
+@pytest.mark.parametrize(
+    "cls,kwargs",
+    [
+        (GaussianNoise, {"std": (-1.0, 1.0)}),
+        (MotionBlur, {"kernel_size": 4}),
+        (JPEGCompression, {"quality": (0, 75)}),
+        (GaussianPatchBrightness, {"sigma_range": (0.0, 0.25)}),
+        (RandomShadow, {"opacity": (0.3, 1.1)}),
+        (CoarseDropout, {"max_holes": 0}),
+        (GammaCorrection, {"gamma": 0.0}),
+        (PlanckianJitter, {"temperature": (2_000, 6_500)}),
+    ],
+)
+def test_robotics_transform_rejects_invalid_config(cls, kwargs):
+    with pytest.raises(ValueError):
+        cls(**kwargs)