jhj0517 commited on
Commit
ebdd2a9
1 Parent(s): c89e57a

Fix ndarray type when there's only one box

Browse files
Files changed (1) hide show
  1. 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) == 1:
162
- return [{"segmentation": mask, "area": place_holder} for mask in masks]
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