import gradio as gr
from meta_segment_anything import SegmentAnything
from color import Color
from color_wheel import ColorWheel

def process_image(image):
	if image is None:
		yield None
		return

	color = Color.fromRgb(0xff, 0x0, 0x0)
	cwt = ColorWheel(color)
	colors = []
	for t in cwt.tone15:
		cw = ColorWheel(Color.fromRgb(t.r, t.g, t.b))
		for h in cw.hue:
			colors.append((h.r, h.g, h.b, 0x80))

	generator = SegmentAnything()
	masks = generator.generate(image)
	masks = sorted(masks, key=lambda x: x['area'], reverse=True)
	index = 0
	for mask in masks:
		if index >= len(colors):
			break
		maskimage = SegmentAnything.makeMaskImage(mask['segmentation'].T, colors[index])
		image.paste(maskimage, (0, 0), maskimage)
		index += 1
		yield image

app = gr.Interface(
	title='Segment Anything Demo',
	description='generate masks for an entire image',
	fn=process_image,
	inputs=gr.Image(type='pil'),
	outputs=gr.Image(label='segmentation', type='pil'),
	allow_flagging='never',
	examples=[['examples/example1.jpg'], ['examples/example2.jpg']],
	concurrency_limit=20,
	#cache_examples=False
)
app.launch()