fix: change inference function

app.py

@@ -2,18 +2,18 @@ from transformers import DetrForObjectDetection, DetrImageProcessor
 import torch
 from PIL import Image
 import matplotlib.pyplot as plt
-import io
 import gradio as gr
-from random import choice
+import cv2
+import torch
+import supervision as sv
+import numpy as np
 
 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 CHECKPOINT = 'facebook/detr-resnet-50'
 CHECKPOINT_ACCIDENT_DETECTION = 'hilmantm/detr-traffic-accident-detection'
 CONFIDENCE_TRESHOLD = 0.5
 IOU_TRESHOLD = 0.8
-COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
-            "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
-            "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
+NMS_TRESHOLD = 0.5
 fdic = {
     "family" : "Impact",
     "style" : "italic",
@@ -26,34 +26,44 @@ image_processor = DetrImageProcessor.from_pretrained(CHECKPOINT)
 model = DetrForObjectDetection.from_pretrained(CHECKPOINT_ACCIDENT_DETECTION)
 model.to(DEVICE)
 
-def get_figure(in_pil_img, in_results):
-    plt.figure(figsize=(16, 16))
-    plt.imshow(in_pil_img)
-    ax = plt.gca()
-
-    for prediction in in_results:
-        selected_color = choice(COLORS)
-
-        x, y = prediction['box']['xmin'], prediction['box']['ymin'],
-        w, h = prediction['box']['xmax'] - prediction['box']['xmin'], prediction['box']['ymax'] - prediction['box']['ymin']
-
-        ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3))
-        ax.text(x, y, f"{prediction['label']}: {round(prediction['score']*100, 1)}%", fontdict=fdic)
-
-    plt.axis("off")
-
-    return plt.gcf()
-
-def inference_from_image(image):
-    results = model(image)
-    figure = get_figure(image, results)
-
-    buf = io.BytesIO()
-    figure.savefig(buf, bbox_inches='tight')
-    buf.seek(0)
-    output_pil_img = Image.open(buf)
-
-    return output_pil_img
+# use this function only for DETR Algorithm
+# def detect_object(model, test_image_path, nms_treshold = 0.5):
+
+def inference_from_image(pil_image):
+  
+  box_annotator = sv.BoxAnnotator()
+  numpy_image = np.array(pil_image)
+  # Convert BGR to RGB if needed (OpenCV uses BGR by default)
+  opencv_image_bgr = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
+  image = cv2.cvtColor(opencv_image_bgr, cv2.COLOR_BGR2RGB)
+
+  # inference
+  with torch.no_grad():
+      # load image and predict
+      inputs = image_processor(images=image, return_tensors='pt').to(DEVICE)
+      outputs = model(**inputs)
+      # post-process
+      target_sizes = torch.tensor([image.shape[:2]]).to(DEVICE)
+      results = image_processor.post_process_object_detection(
+          outputs=outputs,
+          threshold=CONFIDENCE_TRESHOLD,
+          target_sizes=target_sizes
+      )[0]
+
+  if results['scores'].shape[0] != 0 or results['labels'].shape[0] != 0:
+    # annotate
+    detections = sv.Detections.from_transformers(transformers_results=results).with_nms(threshold=NMS_TRESHOLD)
+    labels = [
+        f"{model.config.id2label[class_id]} {confidence:0.2f}"
+        for _, confidence, class_id, _
+        in detections
+    ]
+    frame = box_annotator.annotate(scene=image.copy(), detections=detections, labels=labels)
+    result_image = Image.fromarray(frame)
+    return result_image
+  else:
+    print("No object detected")
+    return None
 
 with gr.Blocks() as demo:
     gr.Markdown(

requirements.txt

@@ -1,2 +1,7 @@
 torch
-transformers[timm]
+transformers[timm]
+supervision==0.3.0
+pytorch-lightning
+roboflow
+timm
+numpy