Implement segment_from_box

Author: constantinpape

micro_sam/segment_from_prompts.py

@@ -12,13 +12,22 @@ def _compute_box(mask, original_size=None):
         coords[1].min(), coords[0].min(),
         coords[1].max() + 1, coords[0].max() + 1,
     ])
-    # FIXME how do we deal with aspect ratios???
+    # TODO how do we deal with aspect ratios???
     if original_size is not None:
         trafo = ResizeLongestSide(max(original_size))
         box = trafo.apply_boxes(box[None], (256, 256)).squeeze()
     return box
 
 
+def _process_box(box, original_size=None):
+    box_processed = box[[1, 0, 3, 2]]
+    # TODO how do we deal with aspect ratios???
+    if original_size is not None:
+        trafo = ResizeLongestSide(max(original_size))
+        box_processed = trafo.apply_boxes(box[None], (256, 256)).squeeze()
+    return box_processed
+
+
 def _compute_logits(mask, eps=1e-3):
 
     def inv_sigmoid(x):
@@ -77,3 +86,17 @@ def segment_from_mask(
         return mask, scores, logits
     else:
         return mask
+
+
+def segment_from_box(
+    predictor, box,
+    image_embeddings=None, i=None, original_size=None,
+    multimask_output=False, return_all=False,
+):
+    if image_embeddings is not None:
+        util.set_precomputed(predictor, image_embeddings, i)
+    mask, scores, logits = predictor.predict(box=_process_box(box, original_size), multimask_output=multimask_output)
+    if return_all:
+        return mask, scores, logits
+    else:
+        return mask

test/test_segment_from_prompts.py

@@ -51,6 +51,16 @@ def test_segment_from_mask(self):
         predicted = segment_from_mask(predictor, mask, use_mask=True, use_box=False)
         self.assertGreater(util.compute_iou(mask, predicted), 0.9)
 
+    def test_segment_from_box(self):
+        from micro_sam.segment_from_prompts import segment_from_box
+
+        mask, image = self._get_input()
+        predictor = self._get_model(image)
+
+        box = np.array([106, 106, 150, 150])
+        predicted = segment_from_box(predictor, box)
+        self.assertGreater(util.compute_iou(mask, predicted), 0.9)
+
 
 if __name__ == "__main__":
     unittest.main()