Spaces:
Running
Running
jhj0517
commited on
Commit
•
ebdd2a9
1
Parent(s):
c89e57a
Fix ndarray type when there's only one box
Browse files- modules/sam_inference.py +4 -2
modules/sam_inference.py
CHANGED
@@ -137,6 +137,8 @@ class SamInference:
|
|
137 |
image = image_prompt_input_data["image"]
|
138 |
image = np.array(image.convert("RGB"))
|
139 |
box = image_prompt_input_data["points"]
|
|
|
|
|
140 |
box = np.array([[x1, y1, x2, y2] for x1, y1, _, x2, y2, _ in box])
|
141 |
|
142 |
predicted_masks, scores, logits = self.predict_image(
|
@@ -158,7 +160,7 @@ class SamInference:
|
|
158 |
masks: np.ndarray
|
159 |
):
|
160 |
place_holder = 0
|
161 |
-
if len(masks)
|
162 |
-
|
163 |
result = [{"segmentation": mask[0], "area": place_holder} for mask in masks]
|
164 |
return result
|
|
|
137 |
image = image_prompt_input_data["image"]
|
138 |
image = np.array(image.convert("RGB"))
|
139 |
box = image_prompt_input_data["points"]
|
140 |
+
if len(box) == 0:
|
141 |
+
return [image], []
|
142 |
box = np.array([[x1, y1, x2, y2] for x1, y1, _, x2, y2, _ in box])
|
143 |
|
144 |
predicted_masks, scores, logits = self.predict_image(
|
|
|
160 |
masks: np.ndarray
|
161 |
):
|
162 |
place_holder = 0
|
163 |
+
if len(masks.shape) <= 3:
|
164 |
+
masks = np.expand_dims(masks, axis=0)
|
165 |
result = [{"segmentation": mask[0], "area": place_holder} for mask in masks]
|
166 |
return result
|