from __future__ import absolute_import, unicode_literals
import cv2
# from myprofile.models import *
from ultralytics import YOLO


def detect_objects_in_roi(roi_frame):
    # fps = 30
    # fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    # out = cv2.VideoWriter("output.mp4", fourcc, fps, (roi_frame.shape[1], roi_frame.shape[0]))  # Output video filename, codec, fps, frame size
    model = YOLO('best.pt')
    results = model(roi_frame)[0]

    classes = [
            'Helmet: yes', 'Glasses: yes', 'Gloves: yes', 'Shoes: yes', 'Person',
            'Helmet: no', 'Glasses: no', 'Gloves: no', 'Shoes: no'
        ]

    detacted_classes = []
    for result in results:
            x1, y1, x2, y2, conf, class_idx = result.boxes.data.tolist()[0]
            class_idx = int(class_idx)
            class_label = classes[class_idx]
            detacted_classes.append(class_label)

    return detacted_classes


def handle():
        
    cap = cv2.VideoCapture("/home/manan/Manan/visonx_demo/visnx-backend/Nettyfy_visnx/myprofile/management/commands/1.mov")

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print(f"Total number of frames: {total_frames}")
    # Read frames from the RTSP stream
    ret, frame = cap.read()

    detactions = []

    skip = True

    # Initialize loop variable
    counter = 0


    while counter < total_frames:
        counter += 1

        # Skip one iteration
        if skip:
            skip = False  # Reset the skip variable
            continue  # Skip the rest of the loop
        # Capture frame-by-frame
        ret, frame = cap.read()
        if ret:

            detected_objects = detect_objects_in_roi(frame)

            for i in detected_objects:
                if i in ['Helmet: no', 'Glasses: no', 'Gloves: no', 'Shoes: no']:
                    detactions.append(i)
        else:
            break

    print(detactions)
   
handle()