Add point postprocessing

Signed-off-by: heyufan1995 <[email protected]>

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,7 +168,8 @@ 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:
@@ -231,6 +232,8 @@ 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.