Update app.py

app.py

@@ -12,12 +12,6 @@ def safe_load_yolo_model(path):
 model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
 model_best = safe_load_yolo_model('./data/best.pt')
 
-# Class names for YOLO model (replace with actual class names used in your YOLO model)
-yolo_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe']
-
-# Class names for best.pt model (assumed classes for crack and pothole)
-best_classes = ['Crack', 'Pothole']
-
 def process_video(video):
     # Open the video using OpenCV
     cap = cv2.VideoCapture(video)
@@ -45,7 +39,8 @@ def process_video(video):
             for box in boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                 class_id = int(box.cls[0])  # Class index for YOLO
-                label = f"YOLO: {yolo_classes[class_id]} - {box.conf[0]:.2f}"  # Map class_id to class name
+                # Use model's built-in class names
+                label = f"YOLO: {model_yolo11.names[class_id]} - {box.conf[0]:.2f}"
                 cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                 cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
 
@@ -55,13 +50,8 @@ def process_video(video):
             for box in boxes:
                 x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                 class_id = int(box.cls[0])  # Class index for best.pt
-                
-                # Check if the class_id is within the range of best_classes
-                if class_id < len(best_classes):
-                    label = f"Best: {best_classes[class_id]} - {box.conf[0]:.2f}"  # Map class_id to specific labels
-                else:
-                    label = f"Best: Unknown Class - {box.conf[0]:.2f}"  # Handle out of range class_id
-
+                # Use model's built-in class names for best.pt
+                label = f"Best: {model_best.names[class_id]} - {box.conf[0]:.2f}"
                 cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
                 cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)