Add iou_predictions and low_res_masks to return

statement in SahiAutoSegmentation class
kadirnar · Nov 4, 2023 · 2dfc287 · 2dfc287
1 parent 8c28fab
commit 2dfc287
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/metaseg/sahi_predictor.py b/metaseg/sahi_predictor.py
@@ -94,7 +94,7 @@ def image_predict(
         if type(input_box[0]) == list:
             input_boxes, new_boxes = multi_boxes(input_box, predictor, read_image)
 
-            masks, _, _ = predictor.predict_torch(
+            masks, iou_predictions, low_res_masks = predictor.predict_torch(
                 point_coords=None,
                 point_labels=None,
                 boxes=new_boxes,
@@ -104,7 +104,7 @@ def image_predict(
         elif type(input_box[0]) == int:
             input_boxes = np.array(input_box)[None, :]
 
-            masks, _, _ = predictor.predict(
+            masks, iou_predictions, low_res_masks = predictor.predict(
                 point_coords=input_point,
                 point_labels=input_label,
                 box=input_boxes,
@@ -126,4 +126,4 @@ def image_predict(
         if show:
             plt.show()
 
-        return masks
+        return masks, iou_predictions, low_res_masks