Thesis-Demo / gradio_test.py
xuan2k's picture
update thesis demo with SAM
cde08ad
raw
history blame
2.35 kB
import gradio as gr
import cv2
import numpy as np
from PIL import Image
block = gr.Blocks(
title="SAM and others",
# theme="shivi/calm_seafoam@>=0.0.1,<1.0.0",
)
colors = [(255, 0, 0), (0, 255, 0)]
markers = [1, 5]
def get_point(img, sel_pix, evt: gr.SelectData):
img = np.array(img, dtype=np.uint8)
sel_pix.append(evt.index)
# draw points
print(sel_pix)
for point in sel_pix:
cv2.drawMarker(img, point, colors[0], markerType=markers[0], markerSize=6, thickness=2)
return Image.fromarray(img).convert("RGB")
def undo_button(orig_img, sel_pix):
temp = orig_img.copy()
temp = np.array(temp, dtype=np.uint8)
if len(sel_pix) != 0:
sel_pix.pop()
for point in sel_pix:
cv2.drawMarker(temp, point, colors[0], markerType=markers[0], markerSize=6, thickness=2)
return Image.fromarray(temp).convert("RGB")
def toggle_button(orig_img, mode):
print(mode)
if mode:
ret = gr.Image(value= orig_img,elem_id="image_upload", type='pil', label="Upload", height=512, tool = "editor")# tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
else:
ret = gr.Image(value = orig_img, elem_id="image_upload", type='pil', label="Upload", height=512, tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
mode = not mode
return ret, mode
def store_img(img):
print("call for store")
return img, [] # when new image is uploaded, `selected_points` should be empty
with block:
selected_points = gr.State([])
original_image = gr.State()
mode = gr.State(True)
input_image = gr.Image(elem_id="image_upload", type='pil', label="Upload", height=512,)# tool = "sketch", brush_color='#00FFFF', mask_opacity=0.6)
undo = gr.Button("undo mode", visible=True)
toggle = gr.Button("toggle mode", visible=True)
input_image.upload(
store_img,
[input_image],
[original_image, selected_points]
)
input_image.select(
get_point,
[input_image, selected_points],
[input_image]
)
undo.click(fn=undo_button, inputs=[original_image, selected_points], outputs=[input_image])
toggle.click(fn=toggle_button, inputs=[original_image, mode], outputs=[input_image, mode])
block.launch()