Update app.py

app.py

@@ -21,6 +21,17 @@ for i, url in enumerate(file_urls):
     else:
         download_file(file_urls[i], f"image_{i}.jpg")
 
+colors = {
+    0: (255, 0, 0),    # Red for class 0
+    1: (0, 255, 0),    # Green for class 1
+    2: (0, 0, 255),    # Blue for class 2
+    3: (255, 255, 0),  # Yellow for class 3
+    4: (255, 0, 255),  # Magenta for class 4
+    5: (0, 255, 255),  # Cyan for class 5
+    6: (128, 0, 0),    # Maroon for class 6
+    7: (0, 128, 0),    # Green (dark) for class 7
+}
+
 model = YOLO('modelbest.pt')
 path  = [['image_0.jpg'], ['image_1.jpg']]
 video_path = [['video.mp4']]
@@ -31,30 +42,27 @@ def show_preds_image(image_path):
     results = outputs[0].cpu().numpy()
 
     for i, det in enumerate(results.boxes.xyxy):
-        # Draw the bounding box
-        cv2.rectangle(
-            image,
-            (int(det[0]), int(det[1])),
-            (int(det[2]), int(det[3])),
-            color=(0, 0, 255),
-            thickness=2,
-            lineType=cv2.LINE_AA
-        )
-        
-        # Get the class label and confidence score
         class_id = int(results.boxes.cls[i])
-        confidence = results.boxes.conf[i]
-        label = f"{model.names[class_id]}: {confidence:.2f}"
+        label = model.names[class_id]
+
+        # Get the bounding box coordinates
+        x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
+        
+        # Draw the bounding box with the specified color
+        color = colors.get(class_id, (0, 0, 255))
+        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2, cv2.LINE_AA)
         
-        # Draw the label
-        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        top_left = (int(det[0]), int(det[1]) - 10)
-        bottom_right = (int(det[0]) + label_size[0], int(det[1]))
-        cv2.rectangle(image, top_left, bottom_right, (0, 0, 255), cv2.FILLED)
-        cv2.putText(image, label, (int(det[0]), int(det[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+        # Calculate text size and position
+        label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.75, 2)
+        text_x = x1 + (x2 - x1) // 2 - label_size[0] // 2
+        text_y = y1 + (y2 - y1) // 2 + label_size[1] // 2
+
+        # Draw the label text
+        cv2.putText(image, label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2, cv2.LINE_AA)
 
     return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
+    
 # def show_preds_image(image_path):
 #     image = cv2.imread(image_path)
 #     outputs = model.predict(source=image_path)
@@ -94,17 +102,29 @@ def show_preds_video(video_path):
             frame_copy = frame.copy()
             outputs = model.predict(source=frame)
             results = outputs[0].cpu().numpy()
+
             for i, det in enumerate(results.boxes.xyxy):
-                cv2.rectangle(
-                    frame_copy,
-                    (int(det[0]), int(det[1])),
-                    (int(det[2]), int(det[3])),
-                    color=(0, 0, 255),
-                    thickness=2,
-                    lineType=cv2.LINE_AA
-                )
+                class_id = int(results.boxes.cls[i])
+                label = model.names[class_id]
+
+                # Get the bounding box coordinates
+                x1, y1, x2, y2 = int(det[0]), int(det[1]), int(det[2]), int(det[3])
+                
+                # Draw the bounding box with the specified color
+                color = colors.get(class_id, (0, 0, 255))
+                cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color, 2, cv2.LINE_AA)
+                
+                # Calculate text size and position
+                label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.75, 2)
+                text_x = x1 + (x2 - x1) // 2 - label_size[0] // 2
+                text_y = y1 + (y2 - y1) // 2 + label_size[1] // 2
+
+                # Draw the label text
+                cv2.putText(frame_copy, label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2, cv2.LINE_AA)
+
             yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
 
+            
 inputs_video = [
     gr.Video(format="mp4", label="Input Video"),
 ]