# Code credit: [EdgeSAM Demo](https://huggingface.co/spaces/chongzhou/EdgeSAM). import torch import gradio as gr import numpy as np from tinysam import sam_model_registry, SamPredictor from PIL import ImageDraw from utils.tools_gradio import fast_process import copy import argparse snapshot_download("merve/tinysam", local_dir="tinysam") model_type = "vit_t" sam = sam_model_registry[model_type](checkpoint="./tinysam/tinysam.pth") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") sam.to(device=device) sam.eval() predictor = SamPredictor(sam) examples = [ ["assets/1.jpg"], ["assets/2.jpg"], ["assets/3.jpg"], ["assets/4.jpeg"], ["assets/5.jpg"], ["assets/6.jpeg"] ] # Description title = "
TinySAM [GitHub]
" description_p = """ # Instructions for point mode 1. Upload an image or click one of the provided examples. 2. Select the point type. 3. Click once or multiple times on the image to indicate the object of interest. 4. The Clear button clears all the points. 5. The Reset button resets both points and the image. """ description_b = """ # Instructions for box mode 1. Upload an image or click one of the provided examples. 2. Click twice on the image (diagonal points of the box). 3. The Clear button clears the box. 4. The Reset button resets both the box and the image. """ css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }" def reset(session_state): session_state['coord_list'] = [] session_state['label_list'] = [] session_state['box_list'] = [] session_state['ori_image'] = None session_state['image_with_prompt'] = None session_state['feature'] = None return None, session_state def reset_all(session_state): session_state['coord_list'] = [] session_state['label_list'] = [] session_state['box_list'] = [] session_state['ori_image'] = None session_state['image_with_prompt'] = None session_state['feature'] = None return None, None, session_state def clear(session_state): session_state['coord_list'] = [] session_state['label_list'] = [] session_state['box_list'] = [] session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image']) return session_state['ori_image'], session_state def on_image_upload( image, session_state, input_size=1024 ): session_state['coord_list'] = [] session_state['label_list'] = [] session_state['box_list'] = [] input_size = int(input_size) w, h = image.size scale = input_size / max(w, h) new_w = int(w * scale) new_h = int(h * scale) image = image.resize((new_w, new_h)) session_state['ori_image'] = copy.deepcopy(image) session_state['image_with_prompt'] = copy.deepcopy(image) print("Image changed") nd_image = np.array(image) session_state['feature'] = None #predictor.set_image(nd_image) return image, session_state def convert_box(xyxy): min_x = min(xyxy[0][0], xyxy[1][0]) max_x = max(xyxy[0][0], xyxy[1][0]) min_y = min(xyxy[0][1], xyxy[1][1]) max_y = max(xyxy[0][1], xyxy[1][1]) xyxy[0][0] = min_x xyxy[1][0] = max_x xyxy[0][1] = min_y xyxy[1][1] = max_y return xyxy def segment_with_points( label, session_state, evt: gr.SelectData, input_size=1024, better_quality=False, withContours=True, use_retina=True, mask_random_color=False, ): x, y = evt.index[0], evt.index[1] point_radius, point_color = 5, (97, 217, 54) if label == "Positive" else (237, 34, 13) session_state['coord_list'].append([x, y]) session_state['label_list'].append(1 if label == "Positive" else 0) print(f"coord_list: {session_state['coord_list']}") print(f"label_list: {session_state['label_list']}") draw = ImageDraw.Draw(session_state['image_with_prompt']) draw.ellipse( [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color, ) image = session_state['image_with_prompt'] coord_np = np.array(session_state['coord_list']) label_np = np.array(session_state['label_list']) masks, scores, logits = predictor.predict( point_coords=coord_np, point_labels=label_np, ) print(f'scores: {scores}') area = masks.sum(axis=(1, 2)) print(f'area: {area}') annotations = np.expand_dims(masks[scores.argmax()], axis=0) seg = fast_process( annotations=annotations, image=image, device=device, scale=(1024 // input_size), better_quality=better_quality, mask_random_color=mask_random_color, bbox=None, use_retina=use_retina, withContours=withContours, ) return seg, session_state def segment_with_box( session_state, evt: gr.SelectData, input_size=1024, better_quality=False, withContours=True, use_retina=True, mask_random_color=False, ): x, y = evt.index[0], evt.index[1] point_radius, point_color, box_outline = 5, (97, 217, 54), 5 box_color = (0, 255, 0) if len(session_state['box_list']) == 0: session_state['box_list'].append([x, y]) elif len(session_state['box_list']) == 1: session_state['box_list'].append([x, y]) elif len(session_state['box_list']) == 2: session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image']) session_state['box_list'] = [[x, y]] print(f"box_list: {session_state['box_list']}") draw = ImageDraw.Draw(session_state['image_with_prompt']) draw.ellipse( [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color, ) image = session_state['image_with_prompt'] if len(session_state['box_list']) == 2: box = convert_box(session_state['box_list']) xy = (box[0][0], box[0][1], box[1][0], box[1][1]) draw.rectangle( xy, outline=box_color, width=box_outline ) box_np = np.array(xy) masks, scores, _ = predictor.predict( point_coords=None, point_labels=None, box=box_np[None, :], ) annotations = np.expand_dims(masks[scores.argmax()], axis=0) seg = fast_process( annotations=annotations, image=image, device=device, scale=(1024 // input_size), better_quality=better_quality, mask_random_color=mask_random_color, bbox=None, use_retina=use_retina, withContours=withContours, ) return seg, session_state return image, session_state img_p = gr.Image(label="Input with points", type="pil") img_b = gr.Image(label="Input with box", type="pil") with gr.Blocks(css=css, title="EdgeSAM") as demo: session_state = gr.State({ 'coord_list': [], 'label_list': [], 'box_list': [], 'ori_image': None, 'image_with_prompt': None, 'feature': None }) with gr.Row(): with gr.Column(scale=1): # Title gr.Markdown(title) with gr.Tab("Point mode") as tab_p: # Images with gr.Row(variant="panel"): with gr.Column(scale=1): img_p.render() with gr.Column(scale=1): with gr.Row(): add_or_remove = gr.Radio( ["Positive", "Negative"], value="Positive", label="Point Type" ) with gr.Column(): clear_btn_p = gr.Button("Clear", variant="secondary") reset_btn_p = gr.Button("Reset", variant="secondary") with gr.Row(): gr.Markdown(description_p) with gr.Row(): with gr.Column(): gr.Markdown("Try some of the examples below ⬇️") gr.Examples( examples=examples, inputs=[img_p, session_state], outputs=[img_p, session_state], examples_per_page=8, fn=on_image_upload, run_on_click=True ) with gr.Tab("Box mode") as tab_b: # Images with gr.Row(variant="panel"): with gr.Column(scale=1): img_b.render() with gr.Row(): with gr.Column(): clear_btn_b = gr.Button("Clear", variant="secondary") reset_btn_b = gr.Button("Reset", variant="secondary") gr.Markdown(description_b) with gr.Row(): with gr.Column(): gr.Markdown("Try some of the examples below ⬇️") gr.Examples( examples=examples, inputs=[img_b, session_state], outputs=[img_b, session_state], examples_per_page=8, fn=on_image_upload, run_on_click=True ) with gr.Row(): with gr.Column(scale=1): gr.Markdown( "
visitors
") img_p.upload(on_image_upload, [img_p, session_state], [img_p, session_state]) img_p.select(segment_with_points, [add_or_remove, session_state], [img_p, session_state]) clear_btn_p.click(clear, [session_state], [img_p, session_state]) reset_btn_p.click(reset, [session_state], [img_p, session_state]) tab_p.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state]) img_b.upload(on_image_upload, [img_b, session_state], [img_b, session_state]) img_b.select(segment_with_box, [session_state], [img_b, session_state]) clear_btn_b.click(clear, [session_state], [img_b, session_state]) reset_btn_b.click(reset, [session_state], [img_b, session_state]) tab_b.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state]) demo.queue() demo.launch()