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
import numpy as np

resnet_101_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
resnet_101_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
resnet_50_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
resnet_50_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')

def process_image(image, margin, model):
	if model=='detr-resnet-101':
		processor = resnet_101_processor
		model = resnet_101_model
	else:
		processor = resnet_50_processor
		model = resnet_50_model

	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 = []
	newlabel = {}
	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()])
		newlabel[model.config.id2label[label.item()]] = 1
		index += 1
		yield [drawImage, gallery, newlabel, ','.join(labels)]

app = gr.Interface(
	title='Object Detection for Image V2',
	fn=process_image,
	inputs=[
		gr.Image(type='pil'),
		gr.Slider(maximum=100, step=1, label='margin'),
		gr.Radio(["detr-resnet-50", "detr-resnet-101"], value="detr-resnet-50", label="Select the model")
	],
	outputs=[
		gr.Image(label='boxes', type='pil'),
		gr.Gallery(label='gallery', columns=8, height=140),
		gr.Label(label='scores'),
		gr.Textbox(label='text'),
	],
	allow_flagging='never',
	examples=[['examples/Wild.jpg', 0], ['examples/Football-Match.jpg', 0]],
	# theme='HaleyCH/HaleyCH_Theme',
	#cache_examples=False
)
app.queue(concurrency_count=20)
app.launch()