RysonFeng
Add app.py
6f49966
raw
history blame
No virus
1.84 kB
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()