jhj0517 commited on
Commit
16bf670
1 Parent(s): 2c719e3

Update image type and add result formatter

Browse files
Files changed (2) hide show
  1. app.py +1 -1
  2. modules/sam_inference.py +13 -5
app.py CHANGED
@@ -37,7 +37,7 @@ class App:
37
  with gr.Row():
38
  with gr.Column(scale=5):
39
  img_input = gr.Image(label="Input image here")
40
- img_input_prompter = ImagePrompter(label="Prompt image with Box & Point",
41
  visible=self.default_mode == BOX_PROMPT_MODE)
42
 
43
  with gr.Column(scale=5):
 
37
  with gr.Row():
38
  with gr.Column(scale=5):
39
  img_input = gr.Image(label="Input image here")
40
+ img_input_prompter = ImagePrompter(label="Prompt image with Box & Point", type='pil',
41
  visible=self.default_mode == BOX_PROMPT_MODE)
42
 
43
  with gr.Column(scale=5):
modules/sam_inference.py CHANGED
@@ -91,9 +91,6 @@ class SamInference:
91
  box=box,
92
  multimask_output=params["multimask_output"],
93
  )
94
- print(f"masks: {masks}")
95
- print(f"scores: {scores}")
96
- print(f"logits: {logits}")
97
  return masks, scores, logits
98
 
99
  def divide_layer(self,
@@ -129,20 +126,31 @@ class SamInference:
129
 
130
  elif input_mode == BOX_PROMPT_MODE:
131
  image = image_prompt_input_data["image"]
 
132
  box = image_prompt_input_data["points"]
 
133
  predict_image_hparams = {
134
  "multimask_output": params[0]
135
  }
136
 
137
- generated_masks, scores, logits = self.predict_image(
138
  image=image,
139
  model_type=model_type,
140
  box=box,
141
  **predict_image_hparams
142
  )
 
143
 
144
  save_psd_with_masks(image, generated_masks, output_path)
145
  mask_combined_image = create_mask_combined_images(image, generated_masks)
146
  gallery = create_mask_gallery(image, generated_masks)
147
 
148
- return [mask_combined_image] + gallery, output_path
 
 
 
 
 
 
 
 
 
91
  box=box,
92
  multimask_output=params["multimask_output"],
93
  )
 
 
 
94
  return masks, scores, logits
95
 
96
  def divide_layer(self,
 
126
 
127
  elif input_mode == BOX_PROMPT_MODE:
128
  image = image_prompt_input_data["image"]
129
+ image = np.array(image.convert("RGB"))
130
  box = image_prompt_input_data["points"]
131
+ box = np.array([[x1, y1, x2, y2] for x1, y1, _, x2, y2, _ in box])
132
  predict_image_hparams = {
133
  "multimask_output": params[0]
134
  }
135
 
136
+ predicted_masks, scores, logits = self.predict_image(
137
  image=image,
138
  model_type=model_type,
139
  box=box,
140
  **predict_image_hparams
141
  )
142
+ generated_masks = self.format_to_auto_result(predicted_masks)
143
 
144
  save_psd_with_masks(image, generated_masks, output_path)
145
  mask_combined_image = create_mask_combined_images(image, generated_masks)
146
  gallery = create_mask_gallery(image, generated_masks)
147
 
148
+ return [mask_combined_image] + gallery, output_path
149
+
150
+ @staticmethod
151
+ def format_to_auto_result(
152
+ masks: np.ndarray
153
+ ):
154
+ place_holder = 0
155
+ result = [{"segmentation": mask, "area": place_holder} for mask in masks]
156
+ return result