瀏覽代碼

Remove small noisy objects from mask

Kentaro Wada 1 年之前
父節點
當前提交
c6956c726a
共有 1 個文件被更改,包括 7 次插入1 次删除
  1. 7 1
      labelme/ai/models/segment_anything.py

+ 7 - 1
labelme/ai/models/segment_anything.py

@@ -5,7 +5,7 @@ import imgviz
 import numpy as np
 import numpy as np
 import onnxruntime
 import onnxruntime
 import PIL.Image
 import PIL.Image
-import skimage.measure
+import skimage
 
 
 from ...logger import logger
 from ...logger import logger
 
 
@@ -175,6 +175,12 @@ def _compute_mask_from_points(
     masks, _, _ = decoder_session.run(None, decoder_inputs)
     masks, _, _ = decoder_session.run(None, decoder_inputs)
     mask = masks[0, 0]  # (1, 1, H, W) -> (H, W)
     mask = masks[0, 0]  # (1, 1, H, W) -> (H, W)
     mask = mask > 0.0
     mask = mask > 0.0
+
+    MIN_SIZE_RATIO = 0.05
+    skimage.morphology.remove_small_objects(
+        mask, min_size=mask.sum() * MIN_SIZE_RATIO, out=mask
+    )
+
     if 0:
     if 0:
         imgviz.io.imsave(
         imgviz.io.imsave(
             "mask.jpg", imgviz.label2rgb(mask, imgviz.rgb2gray(image))
             "mask.jpg", imgviz.label2rgb(mask, imgviz.rgb2gray(image))