tinysam / app_old.py
xinghaochen's picture
Rename app.py to app_old.py
d86ac0d
raw
history blame
2.01 kB
from huggingface_hub import snapshot_download
import gradio as gr
import numpy as np
import torch
import sys
from tinysam import sam_model_registry, SamPredictor
snapshot_download("merve/tinysam", local_dir="tinysam")
model_type = "vit_t"
sam = sam_model_registry[model_type](checkpoint="./tinysam/tinysam.pth")
predictor = SamPredictor(sam)
def infer(img):
if img is None:
gr.Error("Please upload an image and select a point.")
if img["background"] is None:
gr.Error("Please upload an image and select a point.")
# background (original image) layers[0] ( point prompt) composite (total image)
image = img["background"].convert("RGB")
point_prompt = img["layers"][0]
total_image = img["composite"]
predictor.set_image(np.array(image))
# get point prompt
img_arr = np.array(point_prompt)
if not np.any(img_arr):
gr.Error("Please select a point on top of the image.")
else:
nonzero_indices = np.nonzero(img_arr)
img_arr = np.array(point_prompt)
nonzero_indices = np.nonzero(img_arr)
center_x = int(np.mean(nonzero_indices[1]))
center_y = int(np.mean(nonzero_indices[0]))
input_point = np.array([[center_x, center_y]])
input_label = np.array([1])
masks, scores, logits = predictor.predict(
point_coords=input_point,
point_labels=input_label,
)
result_label = [(masks[scores.argmax(), :, :], "mask")]
return image, result_label
with gr.Blocks() as demo:
gr.Markdown("## TinySAM")
gr.Markdown("**[TinySAM](https://arxiv.org/abs/2312.13789) is a framework to distill Segment Anything Model.**")
gr.Markdown("**To try it out, simply upload an image, click the green tick, and then leave a point mark on what you would like to segment using the pencil on Image Editor.**")
with gr.Row():
with gr.Column():
im = gr.ImageEditor(
type="pil"
)
output = gr.AnnotatedImage()
im.change(infer, inputs=im, outputs=output)
demo.launch(debug=True)