Spaces:
Runtime error
Runtime error
File size: 3,560 Bytes
077fc91 5c0b534 7598e8a 077fc91 5c0b534 077fc91 d4233b7 077fc91 f07135c 077fc91 d4233b7 077fc91 4f4f67b 077fc91 d4233b7 077fc91 d4233b7 077fc91 d4233b7 077fc91 7784349 077fc91 b3873d0 077fc91 d4233b7 7299967 8f90f14 077fc91 |
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 |
import os
import gradio as gr
import numpy as np
import cv2
from PIL import Image
import torch
from inference import SegmentPredictor
sam = SegmentPredictor() #service.get_sam(configs.model_type, configs.model_ckpt_path, configs.device)
red = (255,0,0)
blue = (0,0,255)
annos = []
block = gr.Blocks()
with block:
# States
def point_coords_empty():
return []
def point_labels_empty():
return []
raw_image = gr.Image(type='pil', visible=False)
point_coords = gr.State(point_coords_empty)
point_labels = gr.State(point_labels_empty)
masks = gr.State([])
cutout_idx = gr.State(set())
# UI
with gr.Column():
with gr.Row():
input_image = gr.Image(label='Input', height=512, type='pil')
masks_annotated_image = gr.AnnotatedImage(label='Segments', height=512)
cutout_galary = gr.Gallery(label='Cutouts', object_fit='contain', height=512)
with gr.Row():
with gr.Column(scale=1):
with gr.Row():
point_label_radio = gr.Radio(label='Point Label', choices=[1,0], value=1)
text = gr.Textbox(label='Mask Name')
sam_sgmt_everything_btn = gr.Button('Segment Everything!', variant = 'primary')
sam_decode_btn = gr.Button('Predict using points!', variant = 'primary')
reset_btn = gr.Button('Reset')
# components
components = {point_coords, point_labels, raw_image, masks, cutout_idx, input_image,
point_label_radio, text, reset_btn, sam_sgmt_everything_btn,
sam_decode_btn, 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):
print("encoding")
# encode image on upload
sam.encode(input_image)
print("encoding done")
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)
def on_click_sam_dencode_btn(inputs):
print("inferencing")
image = inputs[raw_image]
generated_mask, _, _ = sam.cond_pred(pts=np.array(inputs[point_coords]), lbls=np.array(inputs[point_labels]))
inputs[masks].append((generated_mask, inputs[text]))
return {masks_annotated_image: (image, inputs[masks])}
sam_decode_btn.click(on_click_sam_dencode_btn, components, [masks_annotated_image, masks, cutout_idx], queue=True)
#sam_sgmt_everything_btn.click(on_sam_sgmt_everything_click, components, [masks_annotated_image, masks, cutout_idx], queue=True)
if __name__ == '__main__':
block.queue()
block.launch() |