Untitled

import cv2
from collections import defaultdict
import numpy as np
import supervision as sv
from ultralytics import YOLO

model = YOLO('./best.pt')

cap = cv2.VideoCapture("./PTZ (Bhanga)_DTP_DTP_20240511115957_20240511124916_81886.mp4")

POLYGONS = [
    np.array([[1113, 620], [1301, 645], [1261, 813], [1116, 622]], dtype=np.int32)
]

track_history = defaultdict(lambda: [])

crossed_objects = defaultdict(lambda: [False, False, False, False])

video_info = sv.VideoInfo.from_video_path("./PTZ (Bhanga)_DTP_DTP_20240511115957_20240511124916_81886.mp4")
with sv.VideoSink("output_polygons.mp4", video_info) as sink:
    
    while cap.isOpened():
        success, frame = cap.read()

        if success:
            results = model.track(frame, persist=True, save=True, tracker="bytetrack.yaml")
            print(results[0].boxes)
            if results[0] is not None and results[0].boxes.id is not None:
                print('coming here..................')
                boxes = results[0].boxes.xywh.cpu().numpy()
                track_ids = results[0].boxes.id.int().cpu().tolist()
                class_labels = results[0].names
                annotated_frame = results[0].plot()
                detections = sv.Detections.from_ultralytics(results[0])

                for box, track_id in zip(boxes, track_ids):
                    x, y, w, h = box
                    track = track_history[track_id]
                    track.append((float(x), float(y)))
                    if len(track) > 30:
                        track.pop(0)

                    for idx, polygon in enumerate(POLYGONS):
                        if cv2.pointPolygonTest(polygon, (int(x), int(y)), False) >= 0:
                            if not crossed_objects[track_id][idx]:
                                crossed_objects[track_id][idx] = True
                                print(f'Object ID {track_id} crossed polygon {idx + 1}')

                                class_name = class_labels[detections[0].class_id[0]]

                                cv2.rectangle(annotated_frame, (int(x - w / 2), int(y - h / 2)), (int(x + w / 2), int(y + h / 2)), (0, 255, 0), 2)
                                cv2.putText(annotated_frame, f"ID: {track_id}, Class: {class_name}", (int(x - w / 2), int(y - h / 2) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                for polygon in POLYGONS:
                    cv2.polylines(annotated_frame, [polygon], isClosed=True, color=(0, 255, 0), thickness=2)

                count_text = f"Objects crossed: {sum([any(crossed) for crossed in crossed_objects.values()])}"
                cv2.putText(annotated_frame, count_text, (1000, 400), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

                sink.write_frame(annotated_frame)
            else:
                # If no boxes were detected, write the frame as it is
                annotated_frame = frame.copy()
                for polygon in POLYGONS:
                    cv2.polylines(annotated_frame, [polygon], isClosed=True, color=(0, 255, 0), thickness=2)
                sink.write_frame(annotated_frame)
        else:
            break

cap.release()