import cv2
import numpy as np
from ultralytics import YOLO
import cvzone

device = 'cuda'

# Load the YOLO11 model
model = YOLO("yolo11m-pecan.pt")

# Open the video file (use video file or webcam, here using webcam)
cap = cv2.VideoCapture(0)
cy1=550
offset=60
idDict={}
pecanCount = 0

# Define the desired resolution and FPS
width, height = 1280, 720

# Set the resolution and other properties for each camera
cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)

while True:
    ret,frame = cap.read()
    if not ret:
        break
    
    # Define the black box position (adjust as needed)
    top_left = (0, 0) 
    bottom_right = (1280, 220)  

    # Draw a black rectangle (filled)
    cv2.rectangle(frame, top_left, bottom_right, (0, 0, 0), thickness=-1)

    # Run YOLO11 tracking on the frame, persisting tracks between frames
    results = model.track(frame, persist=True,classes=0,device = device)

    # Check if there are any boxes in the results
    if results[0].boxes is not None and results[0].boxes.id is not None:
        # Get the boxes (x, y, w, h), class IDs, track IDs, and confidences
        boxes = results[0].boxes.xyxy.int().cpu().tolist()  # Bounding boxes
        class_ids = results[0].boxes.cls.int().cpu().tolist()  # Class IDs
        track_ids = results[0].boxes.id.int().cpu().tolist()  # Track IDs
        confidences = results[0].boxes.conf.cpu().tolist()  # Confidence score
       
        for box, class_id, track_id, conf in zip(boxes, class_ids, track_ids, confidences):
            x1, y1, x2, y2 = box
            cy = int(y1+y2)//2
            if track_id in idDict.keys():
                cv2.rectangle(frame,(x1,y1),(x2,y2),(255,255,255),2)
                cvzone.putTextRect(frame,f'{idDict[track_id]}',(x1,y2),1,1)
            else:
                cv2.rectangle(frame,(x1,y1),(x2,y2),(0,255,0),2)
                # cvzone.putTextRect(frame,f'{track_id}',(x1,y2),1,1)
            # cv2.circle(frame,(cx,cy),4,(255,0,0),-1)
            # cvzone.putTextRect(frame,f'{c}',(x1,y1),1,1)
            if cy<(cy1+offset) and cy>(cy1-offset) and track_id not in idDict.keys():
                pecanCount += 1
                idDict[track_id] = pecanCount
            cv2.putText(frame, f'Totals Pecans: {pecanCount}', (50, 50),cv2.FONT_HERSHEY_SIMPLEX, 
                    0.6, [255, 255, 255], 1)
            
    cv2.line(frame,(0,550),(1280,550),(255,0,255),2)
    cv2.line(frame,(0,cy1+offset),(1280,cy1+offset),(0,0,255),2)
    cv2.line(frame,(0,cy1-offset),(1280,cy1-offset),(0,0,255),2)

    cv2.imshow('Frame', frame)
    cv2.imshow("RGB", frame)
    if cv2.waitKey(1) & 0xFF == ord("q"):
       break

# Release the video capture object and close the display window
cap.release()
cv2.destroyAllWindows()
print(pecanCount)