From d71a2d114cf13e41d2382450b0ed68e2cb4a799d Mon Sep 17 00:00:00 2001
From: Yufan He <59374597+heyufan1995@users.noreply.github.com>
Date: Tue, 16 Jul 2024 08:25:35 -0400
Subject: [PATCH] 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 <heyufan1995@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 scripts/infer.py             | 9 +++++++--
 scripts/utils/trans_utils.py | 2 +-
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/scripts/infer.py b/scripts/infer.py
index 9b2be48..f9102e8 100644
--- a/scripts/infer.py
+++ b/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()
diff --git a/scripts/utils/trans_utils.py b/scripts/utils/trans_utils.py
index c446bad..ec3ac4d 100644
--- a/scripts/utils/trans_utils.py
+++ b/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.