import warnings
import time
import threading
import queue
import cv2
from ultralytics import YOLO

warnings.filterwarnings('ignore')

def draw_fps(frame, fps):
    font = cv2.FONT_HERSHEY_SIMPLEX
    cv2.putText(frame, f'FPS: {fps:.2f}', (10, 30), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
    return frame

def capture_frames(cap, frame_queue):
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        if frame_queue.full():
            continue
        frame_queue.put(frame)

if __name__ == '__main__':
    model = YOLO('dyhead_best.pt')  # select your YOLOv8 model.pt path
    
    # Open a video capture (0 for webcam)
    cap = cv2.VideoCapture(0)
    
    # Create a queue to store frames
    frame_queue = queue.Queue(maxsize=1)
    
    # Start the frame capture thread
    capture_thread = threading.Thread(target=capture_frames, args=(cap, frame_queue))
    capture_thread.daemon = True
    capture_thread.start()
    
    prev_detection_time = 0
    fps = 0
    img_size = 640  # Set your desired image size here

    while True:
        if not frame_queue.empty():
            frame = frame_queue.get()
            
            # Run object detection
            detection_start_time = time.time()
            results = model(frame, imgsz=img_size)  # YOLOv8 inference with specified image size
            detection_end_time = time.time()
            
            # Calculate FPS based on detection time
            detection_time = detection_end_time - detection_start_time
            if detection_time > 0:
                fps = 1 / detection_time  # calculation of FPS
            
            # Draw bounding boxes and labels on the frame
            annotated_frame = results[0].plot()
            
            # Draw FPS on the frame
            annotated_frame = draw_fps(annotated_frame, fps)
            
            # Display the frame with OpenCV
            cv2.imshow('YOLOv8 Detection', annotated_frame)
            
            # Break the loop if 'q' is pressed
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
            
            print(f"Detection time: {detection_time:.3f}s, FPS: {fps:.2f}")

    # Release the video capture object and close all OpenCV windows
    cap.release()
    cv2.destroyAllWindows()