feat: adjust model processing

.gitignore

@@ -0,0 +1,3 @@
+.DS_Store
+.idea
+.env

app.py

@@ -1,20 +1,44 @@
 import gradio as gr
 import torch
-from transformers import pipeline
+import cv2
 
+model = torch.hub.load('ultralytics/yolov5', 'custom', 'model/best.onnx')
+
+
+CLASS_COLORS = {
+    0: [148, 0, 211],  # class 1 (violet)
+    1: [255, 0, 0],    # class 2 (red)
+    2: [255, 127, 0],  # class 3 (orange)
+    3: [255, 255, 0],  # class 4 (yellow)
+    4: [0, 255, 0],    # class 5 (green)
+    5: [0, 0, 255],    # class 6 (blue)
+}
 
-model = pipeline('onnx', model='model/best.onnx')
 
 def object_detection(image):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
     results = model(image)
 
-    return {
-        'image': image,
-        'detections': results
-    }
+    bboxes = results.xyxy[0].tolist()
+    labels = results.xyxy[0][:, -1].long().tolist()
+    scores = results.xyxy[0][:, -2].tolist()
+
+    for bbox, label, score in zip(bboxes, labels, scores):
+        label_name = results.names[label]
+        color = CLASS_COLORS[label]
+
+        cv2.rectangle(image, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 2)
+        text = f"{label_name} ({score:.2f})"
+        cv2.putText(image, text, (int(bbox[0]), int(bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    return image
+
 
-inputs = gr.inputs.Image()
-outputs = gr.outputs.Image(label='Input Image', type='numpy'), gr.outputs.JSON(label='Object Detections')
+inputs = gr.components.Image(shape=(640, 640))
+outputs = gr.components.Image(label='Input Image', shape=(640, 640))
 
 iface = gr.Interface(fn=object_detection, inputs=inputs, outputs=outputs)
 

model/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0dd7ed03452e8faaf24a3d3c7b4b6694e369ab54d6c807dc48b5c0d6ac675fc1
+size 14468157

requirements.txt

@@ -1,2 +1,3 @@
+gradio==3.28.3
+opencv-python==4.7.0.72
 torch==2.0.0
-git+https://github.com/huggingface/transformers.git

yolov5s.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b3b748c1e592ddd8868022e8732fde20025197328490623cc16c6f24d0782ee
+size 14808437