app.py

import numpy as np 
import os 
 
from nanosam import Predictor 
 
import gradio as gr 
import time 
from PIL import ImageDraw 
from utils import download_file_from_url, fast_process, format_results, point_prompt 
 
# Most of our demo code is from [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM). Huge thanks for AN-619. 
 
if not os.path.exists("onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx"): 
 download_file_from_url( 
 "https://huggingface.co/dragonSwing/nanosam/resolve/main/sam_hgv2_b4_ln_nonorm_image_encoder.onnx", 
 "onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx", 
 ) 
 
if not os.path.exists("onnx/efficientvit_l0_mask_decoder.onnx"): 
 download_file_from_url( 
 "https://huggingface.co/dragonSwing/nanosam/resolve/main/efficientvit_l0_mask_decoder.onnx", 
 "onnx/efficientvit_l0_mask_decoder.onnx", 
 ) 
 
# Load the pre-trained model 
image_encoder_cfg = { 
 "path": "onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx", 
 "provider": "cpu", 
 "normalize_input": False, 
} 
mask_decoder_cfg = { 
 "path": "onnx/efficientvit_l0_mask_decoder.onnx", 
 "provider": "cpu", 
} 
predictor = Predictor(image_encoder_cfg, mask_decoder_cfg) 
 
# Description 
title = "<center><strong><font size='8'>Faster Segment Anything(NanoSAM)<font></strong></center>" 
 
description_p = """ ## This is a demo of [Faster Segment Anything(NanoSAM) Model](https://github.com/binh234/nanosam). 
 # Instructions for point mode 
 0. Restart by click the Restart button 
 1. Select a point with Add Mask for the foreground (Must) 
 2. Select a point with Remove Area for the background (Optional) 
 3. Click the Start Segmenting. 
 - Github [link](https://github.com/binh234/nanosam) 
 - Model Card [link](https://huggingface.co/dragoswing/nanosam) 
 We will provide box mode soon. 
 Enjoy! 
 """ 
 
examples = [ 
 ["assets/picture3.jpg"], 
 ["assets/picture4.jpg"], 
 ["assets/picture5.jpg"], 
 ["assets/picture6.jpg"], 
 ["assets/picture1.jpg"], 
 ["assets/picture2.jpg"], 
 ["assets/dogs.jpg"], 
] 
 
css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }" 
 
 
def get_empty_state(): 
 return {"points": [], "point_labels": [], "features": None} 
 
 
def clear(): 
 return None, None, get_empty_state() 
 
 
def set_image(image): 
 state = get_empty_state() 
 start = time.perf_counter() 
 predictor.set_image(image) 
 end = time.perf_counter() 
 print(f"Encoder time: {end - start: .3f}s") 
 state["features"] = predictor.features 
 return state 
 
 
def segment_with_points( 
 image, 
 state, 
 better_quality=False, 
 withContours=True, 
 use_retina=True, 
 mask_random_color=True, 
): 
 global predictor 
 
 points = np.asarray(state["points"]) 
 point_labels = np.asarray(state["point_labels"]) 
 if len(points) == 0 and len(point_labels) == 0: 
 raise gr.Error("No points selected") 
 if len(points) != len(point_labels): 
 raise gr.Error("Mismatch length between points and point labels") 
 if state["features"] is None: 
 raise gr.Error( 
 "Image was not set correctly, please wait for a moment after uploading image before drawing points!" 
 ) 
 
 predictor.features = state["features"] 
 img_w, img_h = image.size 
 predictor.original_size = (img_h, img_w) 
 start = time.perf_counter() 
 masks, scores, logits = predictor.predict( 
 points=points, 
 point_labels=point_labels, 
 ) 
 end = time.perf_counter() 
 print(f"Decoder time: {end - start: .3f}s") 
 
 # results = format_results(masks[0], scores[0], logits[0], 0) 
 
 # annotations, _ = point_prompt(results, points, point_labels, img_h, img_w) 
 # annotations = np.array([annotations]) 
 
 fig = fast_process( 
 annotations=[masks[0, scores.argmax()] > 0], 
 image=image, 
 scale=1, 
 better_quality=better_quality, 
 mask_random_color=mask_random_color, 
 bbox=None, 
 use_retina=use_retina, 
 withContours=withContours, 
 ) 
 
 # return fig, None 
 return fig 
 
 
def get_points_with_draw(image, label, evt: gr.SelectData, state): 
 x, y = evt.index[0], evt.index[1] 
 point_radius, point_color = 15, ( 
 (255, 255, 0) 
 if label == "Add Mask" 
 else ( 
 255, 
 0, 
 255, 
 ) 
 ) 
 state["points"].append([x, y]) 
 state["point_labels"].append(1 if label == "Add Mask" else 0) 
 
 print(x, y, label == "Add Mask") 
 
 draw = ImageDraw.Draw(image) 
 draw.ellipse( 
 [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], 
 fill=point_color, 
 ) 
 return image, state 
 
 
cond_img_p = gr.Image(label="Input with points", type="pil", interactive=True) 
 
segm_img_p = gr.Image(label="Segmented Image with points", interactive=False, type="pil") 
 
global_points = [] 
global_point_labels = [] 
 
with gr.Blocks(css=css, title="Faster Segment Anything(NanoSAM)") as demo: 
 state = gr.State(value=get_empty_state()) 
 with gr.Row(): 
 with gr.Column(scale=1): 
 # Title 
 gr.Markdown(title) 
 
 with gr.Tab("Point mode"): 
 # Images 
 with gr.Row(variant="panel"): 
 with gr.Column(scale=1): 
 cond_img_p.render() 
 
 with gr.Column(scale=1): 
 segm_img_p.render() 
 
 # Submit & Clear 
 with gr.Row(): 
 with gr.Column(): 
 with gr.Row(): 
 add_or_remove = gr.Radio( 
 ["Add Mask", "Remove Area"], 
 value="Add Mask", 
 ) 
 
 with gr.Column(): 
 segment_btn_p = gr.Button("Start segmenting!", variant="primary") 
 restart_btn_p = gr.Button("Restart", variant="secondary") 
 
 gr.Markdown("Try some of the examples below ⬇️") 
 gr.Examples( 
 examples=examples, 
 inputs=[cond_img_p], 
 outputs=[state], 
 fn=set_image, 
 run_on_click=True, 
 examples_per_page=4, 
 ) 
 
 with gr.Column(): 
 # Description 
 gr.Markdown(description_p) 
 
 cond_img_p.upload(set_image, inputs=[cond_img_p], outputs=[state]) 
 cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove, state], [cond_img_p, state]) 
 segment_btn_p.click(segment_with_points, [cond_img_p, state], [segm_img_p]) 
 restart_btn_p.click(clear, outputs=[cond_img_p, segm_img_p, state]) 
 
demo.queue().launch()