Update app.py

app.py

@@ -4,37 +4,16 @@ import torch
 
 from PIL import Image
 import requests
-#from transformers import DetrImageProcessor
-#from transformers import DetrForObjectDetection
 from transformers import pipeline
 import matplotlib.pyplot as plt
 import io
 
-
-#processor = DetrImageProcessor.from_pretrained("sergiopaniego/detr-resnet-50-dc5-fashionpedia-finetuned")
-#model = DetrForObjectDetection.from_pretrained("sergiopaniego/detr-resnet-50-dc5-fashionpedia-finetuned")
 model_pipeline = pipeline("object-detection", model="sergiopaniego/detr-resnet-50-dc5-fashionpedia-finetuned")
 
 
 COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
           [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
 
-'''
-def get_output_figure(pil_img, scores, labels, boxes, threshold):
-    plt.figure(figsize=(16, 10))
-    plt.imshow(pil_img)
-    ax = plt.gca()
-    colors = COLORS * 100
-    for score, label, (xmin, ymin, xmax, ymax), c in zip(scores.tolist(), labels.tolist(), boxes.tolist(), colors):
-        if score > threshold:
-            ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=c, linewidth=3))
-            text = f'{model.config.id2label[label]}: {score:0.2f}'
-            ax.text(xmin, ymin, text, fontsize=15, 
-                    bbox=dict(facecolor='yellow', alpha=0.5))
-    plt.axis('off')
-
-    return plt.gcf()
-'''
 
 def get_output_figure(pil_img, results, threshold):
     plt.figure(figsize=(16, 10))
@@ -58,23 +37,10 @@ def get_output_figure(pil_img, results, threshold):
 
 @spaces.GPU
 def detect(image):
-    #encoding = processor(image, return_tensors='pt')
-    #print(encoding.keys())
-    
-    #with torch.no_grad():
-    #    outputs = model(**encoding)
-
-
     results = model_pipeline(image)
     print(results)
 
-    #width, height = image.size
-    #postprocessed_outputs = processor.post_process_object_detection(outputs, target_sizes=[(height, width)], threshold=0.5)
-    #results = postprocessed_outputs[0]
-
-
-    #output_figure = get_output_figure(image, results['scores'], results['labels'], results['boxes'], threshold=0.5)
-    output_figure = get_output_figure(image, results, threshold=0.5)
+    output_figure = get_output_figure(image, results, threshold=0.7)
 
     buf = io.BytesIO()
     output_figure.savefig(buf, bbox_inches='tight')