Update app.py

app.py

@@ -4,12 +4,16 @@ from transformers import DetrImageProcessor, DetrForObjectDetection
 from color import Color
 from color_wheel import ColorWheel
 from PIL import ImageDraw, ImageFont
+import numpy as np
 
-processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
-model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
+def process_image(image, margin, model):
+	if model=='detr-resnet-101':
+		processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
+		model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
+	else:
+		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
@@ -30,6 +34,7 @@ def process_image(image, margin):
 	index = 0
 	gallery = []
 	labels = []
+	newlabel = {}
 	drawImage = image.copy()
 	draw = ImageDraw.Draw(drawImage)
 	for score, label, box in zip(results['scores'], results['labels'], results['boxes']):
@@ -43,8 +48,14 @@ def process_image(image, 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()])
+		print(model.config.id2label[label.item()])
+		newlabel[model.config.id2label[label.item()]] = 1
 		index += 1
-		yield [drawImage, gallery, ','.join(labels)]
+		print('----------')
+		print(labels)
+		print(results['scores'])
+		print(newlabel)
+		yield [drawImage, gallery, newlabel, ','.join(labels)]
 
 app = gr.Interface(
 	title='Object Detection for Image',
@@ -52,14 +63,17 @@ app = gr.Interface(
 	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)