tests.workflow_auto_label_masks_test

  1import fiftyone as fo
  2import pytest
  3from fiftyone.utils.huggingface import load_from_hub
  4import config.config
  5
  6from main import workflow_auto_label_mask
  7from utils.logging import configure_logging
  8
  9@pytest.fixture(scope="session")
 10def dataset_v51():
 11    dataset_name_hub = "Voxel51/fisheye8k"
 12    dataset_name = "fisheye8k_mask_test"
 13    try:
 14        dataset = fo.load_dataset(dataset_name)
 15    except ValueError:
 16        dataset = load_from_hub(
 17            repo_id=dataset_name_hub,
 18            max_samples=1,
 19            name=dataset_name,
 20        )
 21    return dataset
 22
 23@pytest.fixture(autouse=True)
 24def setup_logging():
 25    configure_logging()
 26
 27@pytest.fixture(autouse=True)
 28def deactivate_wandb_sync():
 29    config.config.WANDB_ACTIVE = False
 30
 31@pytest.mark.parametrize("workflow_config", [
 32    {   # Test 1: SAM2 without a prompt
 33        "semantic_segmentation": {
 34            "sam2": {
 35                "prompt_field": None,
 36                "models": ["segment-anything-2.1-hiera-tiny-image-torch"]
 37            }
 38        },
 39        "depth_estimation": {}
 40    },
 41    {   # Test 2: SAM2 with a prompt
 42        "semantic_segmentation": {
 43            "sam2": {
 44                "prompt_field": "detections",
 45                "models": ["segment-anything-2.1-hiera-tiny-image-torch"]
 46            }
 47        },
 48        "depth_estimation": {}
 49    },
 50    {   # Test 3: Depth Estimation
 51        "semantic_segmentation": {},
 52        "depth_estimation": {
 53            "dpt": {
 54                "models": {"Intel/dpt-swinv2-tiny-256"}
 55            }
 56        }
 57    },
 58])
 59
 60def test_auto_label_mask(dataset_v51, workflow_config):
 61
 62    dataset_info = {"name": "fisheye8k_mask_test"}
 63
 64    print("\n[TEST] Starting workflow_auto_label_mask integration test")
 65    print(f"[TEST] Dataset Name: {dataset_v51.name}")
 66    print(f"[TEST] Number of samples: {len(dataset_v51)}")
 67
 68
 69    # Run if valid config
 70    if workflow_config["semantic_segmentation"] or workflow_config["depth_estimation"]:
 71        workflow_auto_label_mask(dataset_v51, dataset_info, workflow_config)
 72        print("[TEST] workflow_auto_label_mask completed successfully!")
 73    else:
 74        print("[TEST] Skipping workflow, running assertion checks on unmodified dataset.")
 75
 76    assert len(dataset_v51) > 0, "The dataset should not be empty after processing"
 77
 78    if workflow_config["semantic_segmentation"]:
 79        prompt_field = workflow_config["semantic_segmentation"]["sam2"]["prompt_field"]
 80    else:
 81        prompt_field = None
 82
 83    expected_depth_field = "pred_de_Intel_dpt_swinv2_tiny_256"
 84    expected_sam_field_noprompt = "pred_ss_segment_anything_2_1_hiera_tiny_image_torch_noprompt"
 85    expected_sam_field_prompt = f"pred_ss_segment_anything_2_1_hiera_tiny_image_torch_prompt_{prompt_field}"
 86
 87    field_gt = "detections"
 88
 89    if prompt_field is None:
 90        field_masks = "pred_ss_segment_anything_2_1_hiera_tiny_image_torch_noprompt"
 91    else:
 92        field_masks = f"pred_ss_segment_anything_2_1_hiera_tiny_image_torch_prompt_{prompt_field}"
 93
 94    classes_gt = set()
 95    classes_sam = set()
 96    n_found_fields = 0
 97
 98    for sample in dataset_v51:
 99        print(f"Fields in sample: {sample.field_names}")
100        if workflow_config["semantic_segmentation"]:
101            bboxes_gt = sample[field_gt]
102
103            try:
104                sam_masks = sample[field_masks]
105            except KeyError:
106                print(f"Field {field_masks} not found in sample {sample}")
107
108            try:
109                if prompt_field is not None:
110                    for bbox in bboxes_gt.detections:
111                        classes_gt.add(bbox.label)
112                    for mask in sam_masks.detections:
113                        classes_sam.add(mask.label)
114                    field = sample[expected_sam_field_prompt]
115                else:
116                    field = sample[expected_sam_field_noprompt]
117                n_found_fields += 1
118
119            except Exception as e:
120                print(f"Error: {e}")
121                pass
122        if workflow_config["depth_estimation"]:
123            try:
124                field = sample[expected_depth_field]
125                n_found_fields += 1
126            except:
127                pass
128
129    assert classes_gt == classes_sam, f"Classes in Ground Truth {classes_gt} and SAM Masks {classes_sam}"
130    print(f"[TEST] Found {n_found_fields} new fields in the dataset")
131    assert n_found_fields > 0, "No new fields were added to the dataset"
132
133    print("Class Distribution in Ground Truth Bounding Boxes:")
134    print(f"\tGround Truth Classes: {classes_gt}")
135    print(f"\tSAM Masks: {classes_sam}")
136
137    print("[TEST] Verified that new fields are present in the dataset.")