from ultralytics import YOLO
import cv2

# Load the YOLO model
model = YOLO("make.pt")

# Define the mapping of class indices to car types
class_map = {
0: 'beige',
1: 'black',
2: 'blue',
3: 'brown',
4: 'gold',
5: 'green',
6: 'grey',
7: 'orange',
8: 'pink', 
9: 'purple',
10: 'red',
11: 'sivler',
12: 'tan',
13: 'white',
14: 'yellow'
}
# Open the video file
video_path = 'DATA\greencar.png'
result = model(video_path)

# cap = cv2.VideoCapture(video_path)

# while True:
#     ret, frame = cap.read()
#     if not ret:
#         break

#     # Perform object detection
#     results = model(frame)
    
#     # Assuming the top prediction is what you're interested in
#     top_prediction_index = results[0].probs.top5[0]  # Index of the highest probability class
#     top_prediction_prob = results[0].probs.top5conf[0].item()  # Highest probability

#     # Get the car type from the class_map
#     # car_type = class_map[top_prediction_index]
#     print('/n')
#     print(f"{class_map[top_prediction_index]}")

# cap.release()
# cv2.destroyAllWindows()