import cv2
from ultralytics import YOLO
import requests
import threading

# Specify the correct path to your trained model weights
model_path = r"C:\Users\tomin\OneDrive\Počítač\School\Project one\AI Tai Lung COLAB\best (3).pt"
model = YOLO(model_path)

# Initialize webcam
cap = cv2.VideoCapture(0)  # Use the default camera

if not cap.isOpened():
    print("Error: Could not open webcam.")
    exit()

# Set camera resolution and FPS
width = 1080  # Lower resolution for faster processing
height = 720
fps = 30
cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
cap.set(cv2.CAP_PROP_FPS, fps)

# Function to send commands to the Raspberry Pi
def send_command_to_raspberry_pi(detection, accuracy):
    url = 'http://192.168.168.168:5000/control-light'  # Ensure this matches the Flask app on the Pi
    payload = {'detection': detection, 'accuracy': accuracy}
    try:
        response = requests.post(url, json=payload)
        if response.status_code == 200:
            print(f"Sent detection: {detection} with accuracy {accuracy}%")
        else:
            print("Failed to send command")
    except requests.exceptions.RequestException as e:
        print(f"Error sending command: {e}")

# Function to perform predictions
def predict(frame, result_container):
    results = model.predict(source=frame, save=False, save_txt=False, save_conf=False)
    result_container.append(results)

# Main loop
def main_loop():
    result_container = []
    frame_id = 0

    try:
        while True:
            ret, frame = cap.read()  # Capture frame-by-frame
            if not ret:
                print("Error: Could not read frame.")
                break

            # Perform prediction on every frame to ensure faster updates
            result_container.clear()  # Clear previous results
            prediction_thread = threading.Thread(target=predict, args=(frame, result_container))
            prediction_thread.start()
            prediction_thread.join()

            if result_container:
                results = result_container.pop(0)

                # Render results on the frame
                result_frame = results[0].plot()  # Correctly render the prediction boxes

                # Check for detection and accuracy
                detected = len(results[0].boxes) > 0  # Assuming results[0].boxes contain the detected objects
                accuracy = None
                detection = None
                if detected:
                    accuracy = round(float(results[0].boxes[0].conf) * 100, 2) if results[0].boxes else 0
                    detection = results[0].boxes[0].cls.item()  # Get the class label (0 for "Other-cat", 1 for "Tai-Lung")
                    detection = "Tai-Lung" if detection == 1 else "Other-cat"
                    send_command_to_raspberry_pi(detection, accuracy)
                else:
                    send_command_to_raspberry_pi("No detection", 0)

            # Display the resulting frame
            cv2.imshow('Webcam', result_frame)

            # Break the loop on 'q' key press
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

            frame_id += 1

    except KeyboardInterrupt:
        print("Interrupted by user")
    finally:
        cap.release()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main_loop()