Merge remote-tracking branch 'origin/vista3d' into vista3d-export

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

Author: borisfom

data/README.md

@@ -45,7 +45,6 @@ train_files, _, dataset_specific_transforms, dataset_specific_transforms_val = \
 
 The following steps are necessary for creating a multi-dataset data loader for model training.
 Step 1 and 2 generate persistent JSON files based on the original dataset (the `image` and `label` pairs; without the additional pseudo label or supervoxel-based label), and only need to be run once when the JSON files don't exist.
-Step 3 is optional for generating overall data analysis stats.
 
 ##### 1. Generate data list JSON file
 ```
@@ -73,34 +72,16 @@ creates a JSON file in a format:
 ```
 
 This step includes a 5-fold cross validation splitting and
-some logic for 80-20 training/testing splitting.
+some logic for 80-20 training/testing splitting. User need to modify the code in make_datalists.py for their own dataset. Meanwhile, the "training_transform" should manually added for each dataset.
 
 The `original_label_dict` corresponds to the original dataset label definitions.
 The `label_dict` modifies `original_label_dict` by simply rephrasing the terms.
 For example in Task06, `cancer` is renamed to `lung tumor`.
 The output of this step is multiple JSON files, each file corresponds
 to one dataset.
 
-
-##### 2. Verify data pairs and generate a global label dictionary
-```
-python -m data.datasets
-```
-
-This script computes a super set of labels from all the dataset JSON files.
-The output of this step is a `jsons/label_dict.json` file,
-representing the global label dictionary mapping, from class names to globally unique class indices (integers).
-
-
-##### 3. Compute class frequencies, data transform utilities
-```
-python -m data.analyzer ...
-```
-
-This file (`data/analyzer.py`) contains useful transforms for reading images
-and labels, converting labels from dataset-specific labels to the global labels
-according to `jsons/label_dict.json`.
-
+##### 2. Add label_dict.json and label_mapping.json
+Add new class indexes to `label_dict.json` and the local to global mapping to `label_mapping.json`. 
 
 ## SupverVoxel Generation
 1. Download the segment anything repo and download the ViT-H weights

scripts/export.py

@@ -30,7 +30,7 @@
 
 from .sliding_window import point_based_window_inferer, sliding_window_inference
 from .train import CONFIG
-from .utils.trans_utils import VistaPostTransform
+from .utils.trans_utils import VistaPostTransform, get_largest_connected_component_point
 from .utils.trt_utils import ExportWrapper, TRTWrapper
 import time
 
@@ -62,6 +62,7 @@ def infer_wrapper(inputs, model, **kwargs):
     outputs = model(input_images=inputs, **kwargs)
     return outputs.transpose(1, 0)
 
+
 class InferClass:
     def __init__(self, config_file="./configs/infer.yaml", **override):
         logging.basicConfig(stream=sys.stdout, level=logging.INFO)
@@ -73,7 +74,6 @@ def __init__(self, config_file="./configs/infer.yaml", **override):
         parser.read_config(config_file_)
         parser.update(pairs=_args)
 
-        # We do not use AMP for export 
         self.amp = parser.get_parsed_content("amp")
         input_channels = parser.get_parsed_content("input_channels")
         patch_size = parser.get_parsed_content("patch_size")
@@ -182,10 +182,14 @@ def infer(
             batch_data = self.batch_data
         else:
             batch_data = self.infer_transforms(image_file)
-            batch_data["label_prompt"] = label_prompt
+            if label_prompt is not None:
+                batch_data["label_prompt"] = label_prompt
             batch_data = list_data_collate([batch_data])
             self.batch_data = batch_data
         if point is not None:
+            if type(point) is list:
+                point = np.array(point)[np.newaxis, ...]
+                point_label = np.array(point_label)[np.newaxis, ...]
             point = self.transform_points(
                 point,
                 np.linalg.inv(batch_data["image"].affine[0])
@@ -245,6 +249,10 @@ def infer(
                             meta=batch_data["image"].meta,
                         )
                 self.prev_mask = batch_data["pred"]
+                if label_prompt is None and point is not None:
+                    batch_data["pred"] = get_largest_connected_component_point(
+                        batch_data["pred"], point_coords=point, point_labels=point_label
+                    )
                 batch_data["image"] = batch_data["image"].to("cpu")
                 batch_data["pred"] = batch_data["pred"].to("cpu")
                 torch.cuda.empty_cache()

scripts/infer.py

@@ -30,7 +30,7 @@
 
 from .sliding_window import point_based_window_inferer, sliding_window_inference
 from .train import CONFIG
-from .utils.trans_utils import VistaPostTransform
+from .utils.trans_utils import VistaPostTransform, get_largest_connected_component_point
 
 rearrange, _ = optional_import("einops", name="rearrange")
 sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
@@ -168,10 +168,14 @@ def infer(
             batch_data = self.batch_data
         else:
             batch_data = self.infer_transforms(image_file)
-            batch_data["label_prompt"] = label_prompt
+            if label_prompt is not None:
+                batch_data["label_prompt"] = label_prompt
             batch_data = list_data_collate([batch_data])
             self.batch_data = batch_data
         if point is not None:
+            if type(point) is list:
+                point = np.array(point)[np.newaxis, ...]
+                point_label = np.array(point_label)[np.newaxis, ...]
             point = self.transform_points(
                 point,
                 np.linalg.inv(batch_data["image"].affine[0])
@@ -231,6 +235,10 @@ def infer(
                             meta=batch_data["image"].meta,
                         )
                 self.prev_mask = batch_data["pred"]
+                if label_prompt is None and point is not None:
+                    batch_data["pred"] = get_largest_connected_component_point(
+                        batch_data["pred"], point_coords=point, point_labels=point_label
+                    )
                 batch_data["image"] = batch_data["image"].to("cpu")
                 batch_data["pred"] = batch_data["pred"].to("cpu")
                 torch.cuda.empty_cache()

scripts/utils/trans_utils.py

@@ -195,7 +195,7 @@ def dilate3d(input_tensor, erosion=3):
 
 
 def get_largest_connected_component_point(
-    img: NdarrayTensor, point_coords=None, point_labels=None, post_idx=3
+    img: NdarrayTensor, point_coords=None, point_labels=None
 ) -> NdarrayTensor:
     """
     Gets the largest connected component mask of an image. img is before post process! And will include NaN values.
@@ -349,8 +349,9 @@ def __call__(
                     pred += 0.5  # inplace mapping to avoid cloning pred
                     for i in range(1, object_num + 1):
                         frac = i + 0.5
-                        pred[pred == frac] = torch.tensor(data["label_prompt"][i - 1]).to(pred.dtype)
-                        # pred[pred == frac] = data["label_prompt"][i - 1].to(pred.dtype)
+                        pred[pred == frac] = torch.tensor(
+                            data["label_prompt"][i - 1]
+                        ).to(pred.dtype)
                     pred[pred == 0.5] = 0.0
                 data[keys] = pred
         return data