import cv2
from ultralytics import YOLO

def main():
    # Set the output video resolution
    output_width, output_height = 1280, 720

    cap = cv2.VideoCapture(r"D:\Mithil nettyfy\PPE kit\Test-Video\cam 1.mp4")
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, output_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, output_height)

    # Create a VideoWriter object to save the output video
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter("pexels-mikael-blomkvist-8964933 (2160p).mp4", fourcc, 30.0, (output_width, output_height))

    model = YOLO('best.onnx')

    while True:
        ret, frame = cap.read()
        results = model(frame)[0]
        classes = [
            'Helmet: yes', 'Glasses: yes', 'Gloves: yes', 'Shoes: yes', 'Person',
            'Helmet: no', 'Glasses: no', 'Gloves: no', 'Shoes: no'
        ]

        persons = []

        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]

            if class_label == "Person":
                # Create a new person object with an empty labels set
                person = {"bbox": (x1, y1, x2, y2), "labels": set()}
                persons.append(person)

        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]

            if class_label != "Person":
                cx, cy = (x1 + x2) / 2, (y1 + y2) / 2  # center of the PPE bbox
                for person in persons:
                    px1, py1, px2, py2 = person["bbox"]
                    if px1 <= cx <= px2 and py1 <= cy <= py2:  # if the center of the PPE bbox is within the Person bbox
                        person["labels"].add(class_label)  # As "labels" is now a set, the same label cannot be added twice
                        break
                
        for person in persons:
            # Draw the bbox
            x1, y1, x2, y2 = person["bbox"]
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0,255, 0), 3)

            # Set the text position just above the bbox
            text_position = (int(x1), int(y1) - 10)

            for label in person["labels"]:
                # Add a black rectangle for padding
                text_width, text_height = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)[0]
                rectangle_bgr = (0, 0, 0)
                (text_width, text_height) = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1.2, thickness=1)[0]
                box_coords = ((text_position[0], text_position[1]), (text_position[0] + text_width + 2, text_position[1] - text_height - 2))
                cv2.rectangle(frame, box_coords[0], box_coords[1], rectangle_bgr, cv2.FILLED)

                # Select the color
                if "yes" in label or "Person" in label:
                    text_color = (0, 255, 0)  # Green color
                else:
                    text_color = (0, 0, 255)  # Red color

                # Add the text
                cv2.putText(frame, label, text_position, cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
                # Update text position for next label
                text_position = (text_position[0], text_position[1] - 30)

        # Resize frame to the desired output resolution
        frame = cv2.resize(frame, (output_width, output_height))
        
        # Save frame to the output video
        out.write(frame)

        cv2.imshow("yolov8", frame)

        if cv2.waitKey(30) == ord('q'):
            break

    cap.release()
    out.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":

    main()



#====================================================================================================================