Spaces:
Runtime error
Runtime error
jhj0517
commited on
Commit
•
16bf670
1
Parent(s):
2c719e3
Update image type and add result formatter
Browse files- app.py +1 -1
- 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 |
-
|
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
|