Spaces:
Runtime error
Runtime error
File size: 4,122 Bytes
0dd537b 50a5f00 0dd537b 50a5f00 0dd537b 50a5f00 0dd537b 50a5f00 cba1a87 0dd537b 50a5f00 0dd537b 50a5f00 0dd537b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import os
import app_configs as configs
import service
import gradio as gr
import numpy as np
import cv2
from PIL import Image
import logging
from huggingface_hub import hf_hub_download
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()
sam = None #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device)
red = (255,0,0)
blue = (0,0,255)
def load_sam_instance():
global sam
if sam is None:
gr.Info('Initialising SAM, hang in there...')
if not os.path.exists(configs.model_ckpt_path):
chkpt_path = hf_hub_download("ybelkada/segment-anything", configs.model_ckpt_path)
else:
chkpt_path = configs.model_ckpt_path
sam = service.get_sam(configs.model_type, chkpt_path, configs.device)
return sam
block = gr.Blocks()
with block:
# states
def point_coords_empty():
return []
def point_labels_empty():
return []
point_coords = gr.State(point_coords_empty)
point_labels = gr.State(point_labels_empty)
raw_image = gr.Image(type='pil', visible=False)
# UI
with gr.Row():
with gr.Column():
input_image = gr.Image(label='Input', height=512, type='pil')
with gr.Row():
point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
reset_btn = gr.Button('Reset')
run_btn = gr.Button('Run', variant = 'primary')
gr.Examples(examples=[['examples/cat-256.png','examples/cat-256.png']],inputs=[input_image, raw_image])
with gr.Column():
with gr.Tab('Cutout'):
cutout_gallery = gr.Gallery()
with gr.Tab('Annotation'):
masks_annotated_image = gr.AnnotatedImage(label='Segments')
# components
components = {point_coords, point_labels, raw_image, input_image, point_label_radio, reset_btn, run_btn, cutout_gallery, masks_annotated_image}
# event - init coords
def on_reset_btn_click(raw_image):
return raw_image, point_coords_empty(), point_labels_empty(), None
reset_btn.click(on_reset_btn_click, [raw_image], [input_image, point_coords, point_labels], queue=False)
def on_input_image_upload(input_image):
return input_image, point_coords_empty(), point_labels_empty(), None
input_image.upload(on_input_image_upload, [input_image], [raw_image, point_coords, point_labels], queue=False)
# event - set coords
def on_input_image_select(input_image, point_coords, point_labels, point_label_radio, evt: gr.SelectData):
x, y = evt.index
color = red if point_label_radio == 0 else blue
img = np.array(input_image)
cv2.circle(img, (x, y), 5, color, -1)
img = Image.fromarray(img)
point_coords.append([x,y])
point_labels.append(point_label_radio)
return img, point_coords, point_labels
input_image.select(on_input_image_select, [input_image, point_coords, point_labels, point_label_radio], [input_image, point_coords, point_labels], queue=False)
# event - inference
def on_run_btn_click(inputs):
sam = load_sam_instance()
image = inputs[raw_image]
if len(inputs[point_coords]) == 0:
if configs.enable_segment_all:
masks, _ = service.predict_all(sam, image)
else:
raise gr.Error('Segment-all disabled, set point label(s) before running')
else:
masks, _ = service.predict_conditioned(sam,
image,
point_coords=np.array(inputs[point_coords]),
point_labels=np.array(inputs[point_labels]))
annotated = (image, [(masks[i], f'Mask {i}') for i in range(len(masks))])
cutouts = [service.cutout(image, mask) for mask in masks]
return cutouts, annotated
run_btn.click(on_run_btn_click, components, [cutout_gallery, masks_annotated_image], queue=True)
if __name__ == '__main__':
block.queue()
block.launch() |