import gradio as gr
import torch
from transformers import DetrImageProcessor, DetrForObjectDetection
from color import Color
from color_wheel import ColorWheel
from PIL import ImageDraw, ImageFont

processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')

def process_image(image, margin):
	if image is None:
		yield [None, None, 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))

	inputs = processor(images=image, return_tensors='pt')
	outputs = model(**inputs)
	target_sizes = torch.tensor([image.size[::-1]])
	results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]

	index = 0
	gallery = []
	labels = []
	drawImage = image.copy()
	draw = ImageDraw.Draw(drawImage)
	for score, label, box in zip(results['scores'], results['labels'], results['boxes']):
		if index >= len(colors):
			break
		box = [round(i) for i in box.tolist()]
		box[0] = max(0, box[0] - margin)
		box[1] = max(0, box[1] - margin)
		box[2] = min(image.width, box[2] + margin)
		box[3] = min(image.height, box[3] + margin)
		draw.rectangle([(box[0], box[1]), (box[2], box[3])], outline=colors[index], width=4)
		gallery.append(image.crop((box[0], box[1], box[2], box[3])))
		labels.append(model.config.id2label[label.item()])
		index += 1
		yield [drawImage, gallery, ','.join(labels)]

app = gr.Interface(
	title='Object Detection for Image',
	fn=process_image,
	inputs=[
		gr.Image(type='pil'),
		gr.Slider(maximum=100, step=1, label='margin'),
	],
	outputs=[
		gr.Image(label='boxes', type='pil'),
		gr.Gallery(label='gallery', columns=8, height=140),
		gr.Textbox(label='text'),
	],
	allow_flagging='never',
	examples=[['examples/Wild.jpg', 0], ['examples/Football-Match.jpg', 0]],
	#cache_examples=False
)
app.queue(concurrency_count=20)
app.launch()