beverage model

app.py

@@ -26,7 +26,7 @@ for i, url in enumerate(file_urls):
             f"image_{i}.jpg"
         )
 
-model = YOLO('best.pt')
+model = YOLO('yolov8n.pt')
 path  = [['image_0.jpg'], ['image_1.jpg']]
 video_path = [['video.mp4']]
 
@@ -35,7 +35,9 @@ def show_preds_image(image_path):
     image = cv2.imread(image_path)
     outputs = model.predict(source=image_path)
     results = outputs[0].cpu().numpy()
+    
     for i, det in enumerate(results.boxes.xyxy):
+        # Draw bounding boxes
         cv2.rectangle(
             image,
             (int(det[0]), int(det[1])),
@@ -44,6 +46,26 @@ def show_preds_image(image_path):
             thickness=2,
             lineType=cv2.LINE_AA
         )
+        
+        # Get the class of the object
+        class_id = int(results.boxes.cls[i])
+        class_name = model.names[class_id]  # Assuming model.names has class names
+        label = f"{class_name}"
+        
+        # Display the class label above the bounding box
+        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
+        label_y = max(int(det[1]), label_size[1] + 10)
+        cv2.putText(
+            image,
+            label,
+            (int(det[0]), label_y - 5),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (0, 255, 0),
+            2,
+            lineType=cv2.LINE_AA
+        )
+        
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  
 inputs_image = [