app.py

import gradio as gr
import numpy as np
from meta_segment_anything import SegmentAnything
from PIL import Image, ImageDraw

def check_location(image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5, area):
	if image is None:
		yield None
		return
	if not enable1 and not enable2 and not enable3 and not enable4 and not enable5:
		yield None
		return

	points = []
	if enable1:
		points.append([left1, top1])
	if enable2:
		points.append([left2, top2])
	if enable3:
		points.append([left3, top3])
	if enable4:
		points.append([left4, top4])
	if enable5:
		points.append([left5, top5])
	for point in points:
		left, top = point
		draw = ImageDraw.Draw(image)
		draw.ellipse([(left - 2, top - 2), (left + 3, top + 3)], fill=(255, 0, 0))
		yield image

def process_image(image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5, area):
	if image is None:
		yield None
		return
	if not enable1 and not enable2 and not enable3 and not enable4 and not enable5:
		yield None
		return

	predictor = SegmentAnything()
	points = []
	if enable1:
		points.append([left1, top1])
	if enable2:
		points.append([left2, top2])
	if enable3:
		points.append([left3, top3])
	if enable4:
		points.append([left4, top4])
	if enable5:
		points.append([left5, top5])
	newImage = Image.new('RGBA', image.size)
	for point in points:
		point_coords = np.array([[0, 0], point])
		point_labels = np.array([0, 1])
		masks, _, _ = predictor.predict(image, point_coords, point_labels)
		if area == 'small':
			mask = masks[0]
		elif area == 'medium':
			mask = masks[1]
		else:
			mask = masks[-1]
		maskimage = SegmentAnything.makeMaskImage(mask.T, (0xff, 0xff, 0xff, 0xff))
		maskNewImage = SegmentAnything.makeNewImage(image, maskimage)
		newImage.paste(maskNewImage, (0, 0), maskNewImage)
		yield newImage

def tab_select(evt: gr.SelectData, state):
	state['active'] = evt.index
	print('select {}'.format(evt.index))
	return state

def image_select(evt: gr.SelectData, state, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5):
	if state['active'] == 4:
		return [enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, True, evt.index[0], evt.index[1]]
	if state['active'] == 3:
		return [enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, True, evt.index[0], evt.index[1], enable5, left5, top5]
	if state['active'] == 2:
		return [enable1, left1, top1, enable2, left2, top2, True, evt.index[0], evt.index[1], enable4, left4, top4, enable5, left5, top5]
	elif state['active'] == 1:
		return [enable1, left1, top1, True, evt.index[0], evt.index[1], enable3, left3, top3, enable4, left4, top4, enable5, left5, top5]
	return [True, evt.index[0], evt.index[1], enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5]

with gr.Blocks(title='clip-image') as app:
	state = gr.State({ 'active': 0 })

	gr.Markdown('''
	# Clip Image
	clip an image from given points
	''')
	with gr.Row():
		with gr.Column():
			image = gr.Image(type='pil')
			gr.Markdown('click on the image to position')
			with gr.Tab('point1') as tab1:
				enable1 = gr.Checkbox(label='enable', value=True)
				left1 = gr.Slider(maximum=4000, step=1, label='left')
				top1 = gr.Slider(maximum=4000, step=1, label='top')
			with gr.Tab('point2') as tab2:
				enable2 = gr.Checkbox(label='enable')
				left2 = gr.Slider(maximum=4000, step=1, label='left')
				top2 = gr.Slider(maximum=4000, step=1, label='top')
			with gr.Tab('point3') as tab3:
				enable3 = gr.Checkbox(label='enable')
				left3 = gr.Slider(maximum=4000, step=1, label='left')
				top3 = gr.Slider(maximum=4000, step=1, label='top')
			with gr.Tab('point4') as tab4:
				enable4 = gr.Checkbox(label='enable')
				left4 = gr.Slider(maximum=4000, step=1, label='left')
				top4 = gr.Slider(maximum=4000, step=1, label='top')
			with gr.Tab('point5') as tab5:
				enable5 = gr.Checkbox(label='enable')
				left5 = gr.Slider(maximum=4000, step=1, label='left')
				top5 = gr.Slider(maximum=4000, step=1, label='top')
			btnloc = gr.Button(value='check location')
			area = gr.Dropdown(label='target area', choices=['large', 'medium', 'small'], value='large')
			with gr.Row():
				with gr.Column(min_width=160):
					clearBtn = gr.ClearButton()
				with gr.Column(min_width=160):
					btn = gr.Button(value='Submit')
			inputs = [image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5, area]
		with gr.Column():
			outputs = [gr.Image(label='segmentation', type='pil')]
	tab1.select(tab_select, inputs=state, outputs=state)
	tab2.select(tab_select, inputs=state, outputs=state)
	tab3.select(tab_select, inputs=state, outputs=state)
	tab4.select(tab_select, inputs=state, outputs=state)
	tab5.select(tab_select, inputs=state, outputs=state)
	image.select(image_select, inputs=[state, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5], outputs=[enable1, left1, top1, enable2, left2, top2, enable3, left3, top3, enable4, left4, top4, enable5, left5, top5])
	btnloc.click(check_location, inputs=inputs, outputs=outputs)
	clearBtn.add(inputs + outputs)
	btn.click(process_image, inputs=inputs, outputs=outputs, concurrency_limit=20)

	gr.Examples(
		[['examples/example1.jpg', True, 200, 250, True, 340, 250, False, 0, 0, False, 0, 0, False, 0, 0, 'large'], ['examples/example2.jpg', True, 256, 256, False, 0, 0, False, 0, 0, False, 0, 0, False, 0, 0, 'large']],
		inputs,
		outputs,
		process_image,
		cache_examples=True,
	)

app.launch()