Untitled

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

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

cap = cv2.VideoCapture("rtsp://admin:Mawa0304@119.148.25.122:554/Streaming/channels/202")

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])

while cap.isOpened():
    success, frame = cap.read()

    if success:
        results = model.track(frame, persist=True, save=True, tracker="bytetrack.yaml")
        if results[0] is not None and results[0].boxes.id is not None:
            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 = results[0].boxes

            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.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)

            cv2.imshow('Annotated Frame', annotated_frame)
        else:
            annotated_frame = frame.copy()
            for polygon in POLYGONS:
                cv2.polylines(annotated_frame, [polygon], isClosed=True, color=(0, 255, 0), thickness=2)
            cv2.imshow('Annotated Frame', annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()