Update app.py

app.py

@@ -1,40 +1,84 @@
-from transformers import DetrImageProcessor, DetrForObjectDetection
-from transformers import BlipProcessor, BlipForConditionalGeneration
-import torch
-from PIL import Image
-import requests
 import gradio as gr
+import pandas as pd
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from PIL import Image, ImageDraw
+import torch
+from transformers import DetrImageProcessor, DetrForObjectDetection
 
-box_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-box_model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
 
-caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+#image_processor = AutoImageProcessor.from_pretrained('hustvl/yolos-small')
+#model = AutoModelForObjectDetection.from_pretrained('hustvl/yolos-small')
+image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
+model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", revision="no_timm")
 
-def predict_bounding_boxes(imageurl:str):
-    try:
-        response = requests.get(imageurl, stream=True)
-        response.raise_for_status() 
+colors = ["red",
+          "orange",
+          "yellow",
+          "green",
+          "blue",
+          "indigo",
+          "violet",
+          "brown",
+          "black",
+          "slategray",
+         ]
 
-        image_data = Image.open(response.raw)
-        inputs = box_processor(images=image_data, return_tensors="pt")
-        outputs = box_model(**inputs)
+# Resized image width
+WIDTH = 900
 
-        target_sizes = torch.tensor([image_data.size[::-1]])
-        results = box_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.70)[0]
+def detect(image):
+    print(image)
+    width, height = image.size
+    ratio = float(WIDTH) / float(width)
+    new_h = height * ratio
 
-        detections = [{"score": score.item(), "label": box_model.config.id2label[label.item()], "box": box.tolist()} for score, label, box in zip(results["scores"], results["labels"], results["boxes"])]
+    image = image.resize((int(WIDTH), int(new_h)), Image.Resampling.LANCZOS)
+    
+    inputs = image_processor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+
+    # convert outputs to COCO API
+    target_sizes = torch.tensor([image.size[::-1]])
+    results = image_processor.post_process_object_detection(outputs,threshold=0.9, target_sizes=target_sizes)[0]
 
-        raw_image = image_data.convert('RGB')
-        inputs = caption_processor(raw_image, return_tensors="pt")
-        out = caption_model.generate(**inputs)
-        label = caption_processor.decode(out[0], skip_special_tokens=True)
-        return {"image label": label, "detections": detections}
+    draw = ImageDraw.Draw(image)
     
-    except Exception as e:
-        
-        return {"error": str(e)}
+    # label and the count
+    counts = {}
+
+    for score, label in zip(results["scores"], results["labels"]):
+        label_name = model.config.id2label[label.item()]
+        if label_name not in counts:
+            counts[label_name] = 0
+        counts[label_name] += 1
+
+    count_results = {k: v for k, v in (sorted(counts.items(), key=lambda item: item[1], reverse=True)[:10])}
+    label2color = {}
+    for idx, label in enumerate(count_results):
+        label2color[label] = colors[idx]
+
+    for label, box in zip(results["labels"], results["boxes"]):
+        label_name = model.config.id2label[label.item()]
+
+        if label_name in count_results:
+            box = [round(i, 4) for i in box.tolist()]
+            x1, y1, x2, y2 = tuple(box)
+            draw.rectangle((x1, y1, x2, y2), outline=label2color[label_name], width=2)
+            draw.text((x1, y1), label_name, fill="white")
+
+    df = pd.DataFrame({
+        'label': [label for label in count_results],
+        'counts': [counts[label] for label in count_results]
+    })
+    
+    return image, df, count_results
+
+demo = gr.Interface(
+    fn=detect,
+    inputs=[gr.Image(label="Input image", type="pil")],
+    outputs=[gr.Image(label="Output image"), gr.BarPlot(show_label=False, x="label", y="counts", x_title="Labels", y_title="Counts", vertical=False), gr.Textbox(show_label=False)],
+    title="FB Object Detection",
+    cache_examples=False
+)
 
-app = gr.Interface(fn=predict_bounding_boxes, inputs="text", outputs="json")
-app.api = True
-app.launch()
+demo.launch()