import cv2
from ultralytics import YOLO
import socket
import threading
import time

server_address = ("Replace with the Raspberry Pi's IP address", 8500)  

def setup_socket_client():
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect(server_address)
    return client_socket

def capture_frames(client_socket):
    model_path = 'Path/to/your/best.pt'
    model = YOLO(model_path)
    vid = cv2.VideoCapture(0)

    try:
        while True:
            ret, frame = vid.read()
            if not ret:
                break

            frame = cv2.resize(frame, (640, 480))
            results = model(frame)

            left_count = 0
            right_count = 0

            frame_height, frame_width = frame.shape[:2]
            mid_x = frame_width // 2

            for result in results:
                if result.boxes:
                    for box in result.boxes:
                        x1, y1, x2, y2 = map(int, box.xyxy[0])
                        conf = box.conf.item()
                        cls = box.cls.item()
                        obj_name = model.names[int(cls)]

                        if obj_name == 'toy cars':
                            if (x1 + x2) / 2 < mid_x:
                                left_count += 1
                            else:
                                right_count += 1

                            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            label = f'{obj_name} {conf:.2f}'
                            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

            cv2.imshow('frame', frame)

            try:
                message = f'{left_count},{right_count}'
                client_socket.sendall(message.encode('utf-8'))
            except Exception as e:
                print("Error sending data to Raspberry Pi:", e)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    except KeyboardInterrupt:
        print("\nProgram terminated by user.")
    except Exception as e:
        print("An error occurred:", e)
    finally:
        vid.release()
        cv2.destroyAllWindows()

def main():
    client_socket = setup_socket_client()

    try:
        capture_frames(client_socket)
    finally:
        client_socket.close()

if __name__ == "__main__":
    main()