Add point postprocessing (#27)

Fixes # .

### Description

A few sentences describing the changes proposed in this pull request.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] In-line docstrings updated.

---------

Signed-off-by: heyufan1995 <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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,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.