Update test_full.py

mlcommons · Jul 19, 2023 · cc0a900 · cc0a900
1 parent 16c9a53
commit cc0a900
Showing 1 changed file with 30 additions and 57 deletions.
diff --git a/testing/test_full.py b/testing/test_full.py
@@ -1810,63 +1810,36 @@ def test_generic_augmentation_functions():
     output_tensor = temp(input_tensor)
     assert output_tensor != None, "RGB Augmentation should work"
 
-    # Testing HED transform
-    input_tensor = torch.rand(3, 128, 128, 1)
-    params = {
-        "data_augmentation": {
-            "hed_transform": {},
-            "hed_transform_light": {},
-            "hed_transform_heavy": {},
-        }
-    }
-    temp = global_augs_dict["hed_transform"](
-        params_all_preprocessing_and_augs["data_augmentation"]["hed_transform"]
-    )
-    ranges = [
-        "haematoxylin_bias_range",
-        "eosin_bias_range",
-        "dab_bias_range",
-        "haematoxylin_sigma_range",
-        "eosin_sigma_range",
-        "dab_sigma_range",
-    ]
-
-    default_range = [-0.1, 0.1]
-    for key in ranges:
-        params["data_augmentation"]["hed_transform"].setdefault(key, default_range)
-
-    params["data_augmentation"]["hed_transform"].setdefault(
-        "cutoff_range", [0.05, 0.95]
-    )
-
-    # Check if the params are correctly set for each augmentation type
-    assert params["data_augmentation"]["hed_transform"] == {
-        "haematoxylin_bias_range": [-0.1, 0.1],
-        "eosin_bias_range": [-0.1, 0.1],
-        "dab_bias_range": [-0.1, 0.1],
-        "haematoxylin_sigma_range": [-0.1, 0.1],
-        "eosin_sigma_range": [-0.1, 0.1],
-        "dab_sigma_range": [-0.1, 0.1],
-        "cutoff_range": [0.05, 0.95],
-    }
-
-    # this is for all other augmentations
-    input_tensor = torch.rand(3, 128, 128, 128)
-    for aug in params_all_preprocessing_and_augs["data_augmentation"]:
-        aug_lower = aug.lower()
-        output_tensor = None
-        if aug_lower in global_augs_dict:
-            output_tensor = global_augs_dict[aug](
-                params_all_preprocessing_and_augs["data_augmentation"][aug_lower]
-            )(input_tensor)
-            assert output_tensor != None, "Augmentation should work"
-
-    temp = global_augs_dict["hed_transform"](
-        params_all_preprocessing_and_augs["data_augmentation"]["hed_transform"]
-    )
-    output_tensor = None
-    output_tensor = temp(input_tensor)
-    assert output_tensor != None, "HED Augmentation should work"
+    # Added HED augmentation in gandlf
+    augmenter = HedColorAugmenter(
+        haematoxylin_sigma_range=(-0.1, 0.1),
+        haematoxylin_bias_range=(-0.1, 0.1),
+        eosin_sigma_range=(-0.1, 0.1),
+        eosin_bias_range=(-0.1, 0.1),
+        dab_sigma_range=(-0.1, 0.1),
+        dab_bias_range=(-0.1, 0.1),
+        cutoff_range=(0, 1),
+    )
+
+    assert augmenter._sigma_ranges == [(-0.1, 0.1), (-0.1, 0.1), (-0.1, 0.1)]
+    assert augmenter._bias_ranges == [(-0.1, 0.1), (-0.1, 0.1), (-0.1, 0.1)]
+    assert augmenter._cutoff_range == [0, 1]
+    augmenter.randomize()
+    assert all(-0.1 <= sigma <= 0.1 for sigma in augmenter._sigmas)
+    assert all(-0.1 <= bias <= 0.1 for bias in augmenter._biases)
+    # Create a mock patch
+    patch = np.ones((64, 64, 3), dtype=np.uint8) * 127
+    transformed_patch = augmenter.transform(patch)
+
+    # The output patch should have the same shape as the input patch
+    assert transformed_patch.shape == patch.shape
+    # Create a mock Subject
+    image_data = np.ones((64, 64, 3), dtype=np.uint8) * 127
+    subject = Subject(t1=Image(tensor=image_data))
+
+    transformed_subject = transform(subject)
+
+    assert isinstance(transformed_subject, Subject)
 
     # additional test for elastic
     params_elastic = params_all_preprocessing_and_augs["data_augmentation"]["elastic"]