from ultralytics import YOLO
import cv2
import numpy as np
from ultralytics.utils.plotting import Annotator

# Load YOLO model
model = YOLO('best.pt')

# Open the video file windows
# cap = cv2.VideoCapture('D:\_PVS\Video\C0012.MP4')

# open video file linux
cap = cv2.VideoCapture('/home/boaz/Videos/C0012.MP4')


# Set the desired resolution (720x1280)
new_width = 1280
new_height = 720

# Load the image mask
mask_image = cv2.imread('mask_v2.png', cv2.IMREAD_GRAYSCALE)  # image must be black and white
resized_mask = cv2.resize(mask_image, (new_width, new_height))
binary_mask = resized_mask // 255  # Convert to binary mask (0s and 1s)

# Initialize car count
car_count = 0

# Initialize frame counter
frame_counter = 0

while True:
    _, frame = cap.read()

    if frame is None:
        break

    # Increment the frame counter
    frame_counter += 1

    # Process every 24th frame
    if frame_counter % 240 != 0:
        continue

    # Resize the frame to the desired resolution
    resized_frame = cv2.resize(frame, (new_width, new_height))  # reduce it to 1280 x 720 may be better to keep it
    # higher

    # Apply the binary mask to the resized frame
    masked_frame = cv2.bitwise_and(resized_frame, resized_frame, mask=binary_mask)

    # Convert to RGB
    img = cv2.cvtColor(masked_frame, cv2.COLOR_BGR2RGB)

    # Perform object detection
    results = model.predict(masked_frame, conf=0.4, iou=0.7)

    for r in results:
        boxes = r.boxes
        for box in boxes:
            car_count += 1
            b = box.xyxy[0]
            annotator = Annotator(resized_frame)
            annotator.box_label(b)

    frame_with_boxes = annotator.result()
    cv2.imshow('YOLO V8 Detection', frame_with_boxes)

    # Calculate remaining count
    remaining_count = 75 - car_count

    # Print the remaining count
    print(f'Remaining Count: {remaining_count}')

    # Reset car count
    car_count = 0

    if cv2.waitKey(1) & 0xFF == ord(' '):  # if you press the space-bar key, it stops the program
        break

cap.release()
cv2.destroyAllWindows()