Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
|
@@ -18,11 +18,9 @@ import torch
|
|
| 18 |
from segment_anything import sam_model_registry, SamPredictor
|
| 19 |
|
| 20 |
# --- CONFIG ---
|
| 21 |
-
SAM_CHECKPOINT = "sam_vit_h_4b8939.pth"
|
| 22 |
SAM_MODEL_TYPE = "vit_h"
|
| 23 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 24 |
BLUR_RADIUS = 10
|
| 25 |
-
DEMO_IMAGE = "demo.png" # put a demo image in the repo
|
| 26 |
# --------------
|
| 27 |
|
| 28 |
sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
|
|
@@ -64,14 +62,20 @@ def isolate_with_click(image: Image.Image, evt: gr.SelectData):
|
|
| 64 |
rgba = np.dstack((fg_rgb, (fg_alpha * 255).astype(np.uint8)))
|
| 65 |
return Image.fromarray(rgba)
|
| 66 |
|
|
|
|
| 67 |
with gr.Blocks() as demo:
|
| 68 |
-
gr.Markdown("### SAM Object Isolation\nUpload an image
|
| 69 |
|
| 70 |
inp = gr.Image(type="pil", label="Upload image", interactive=True)
|
| 71 |
-
demo_img = gr.Image(value=DEMO_IMAGE, type="pil", label="Demo image (click to try)", interactive=True)
|
| 72 |
out = gr.Image(type="pil", label="Isolated cutout (RGBA)")
|
| 73 |
|
| 74 |
inp.select(isolate_with_click, inputs=[inp], outputs=out)
|
| 75 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 76 |
|
| 77 |
demo.launch()
|
|
|
|
| 18 |
from segment_anything import sam_model_registry, SamPredictor
|
| 19 |
|
| 20 |
# --- CONFIG ---
|
|
|
|
| 21 |
SAM_MODEL_TYPE = "vit_h"
|
| 22 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 23 |
BLUR_RADIUS = 10
|
|
|
|
| 24 |
# --------------
|
| 25 |
|
| 26 |
sam = sam_model_registry[SAM_MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
|
|
|
|
| 62 |
rgba = np.dstack((fg_rgb, (fg_alpha * 255).astype(np.uint8)))
|
| 63 |
return Image.fromarray(rgba)
|
| 64 |
|
| 65 |
+
# --- Gradio UI ---
|
| 66 |
with gr.Blocks() as demo:
|
| 67 |
+
gr.Markdown("### SAM Object Isolation\nUpload an image, then click on the object to isolate it.")
|
| 68 |
|
| 69 |
inp = gr.Image(type="pil", label="Upload image", interactive=True)
|
|
|
|
| 70 |
out = gr.Image(type="pil", label="Isolated cutout (RGBA)")
|
| 71 |
|
| 72 |
inp.select(isolate_with_click, inputs=[inp], outputs=out)
|
| 73 |
+
|
| 74 |
+
# Examples section at the bottom
|
| 75 |
+
gr.Examples(
|
| 76 |
+
examples=["demo.jpg"], # put demo.jpg in repo
|
| 77 |
+
inputs=inp,
|
| 78 |
+
label="Try with demo image"
|
| 79 |
+
)
|
| 80 |
|
| 81 |
demo.launch()
|