gradio/demo/image_selections/run.py
aliabid94 379b6f662e
Get Image click coordinates from .select (#3786)
* changes

* changes

* fix demo notebook

* changes

* changes

* chagnes

* changes

* changes

---------

Co-authored-by: Abubakar Abid <abubakar@huggingface.co>
2023-04-10 14:47:33 -07:00

46 lines
1.7 KiB
Python

import gradio as gr
import numpy as np
with gr.Blocks() as demo:
tolerance = gr.Slider(label="Tolerance", info="How different colors can be in a segment.", minimum=0, maximum=256*3, value=50)
with gr.Row():
input_img = gr.Image(label="Input")
output_img = gr.Image(label="Selected Segment")
def get_select_coords(img, tolerance, evt: gr.SelectData):
visited_pixels = set()
pixels_in_queue = set()
pixels_in_segment = set()
start_pixel = img[evt.index[1], evt.index[0]]
pixels_in_queue.add((evt.index[1], evt.index[0]))
while len(pixels_in_queue) > 0:
pixel = pixels_in_queue.pop()
visited_pixels.add(pixel)
neighbors = []
if pixel[0] > 0:
neighbors.append((pixel[0] - 1, pixel[1]))
if pixel[0] < img.shape[0] - 1:
neighbors.append((pixel[0] + 1, pixel[1]))
if pixel[1] > 0:
neighbors.append((pixel[0], pixel[1] - 1))
if pixel[1] < img.shape[1] - 1:
neighbors.append((pixel[0], pixel[1] + 1))
for neighbor in neighbors:
if neighbor in visited_pixels:
continue
neighbor_pixel = img[neighbor[0], neighbor[1]]
if np.abs(neighbor_pixel - start_pixel).sum() < tolerance:
pixels_in_queue.add(neighbor)
pixels_in_segment.add(neighbor)
out = img.copy() * 0.2
out = out.astype(np.uint8)
for pixel in pixels_in_segment:
out[pixel[0], pixel[1]] = img[pixel[0], pixel[1]]
return out
input_img.select(get_select_coords, [input_img, tolerance], output_img)
if __name__ == "__main__":
demo.launch()