Untitled

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

torch.cuda.set_device(0) # Set to your desired GPU number

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

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

POLYGONS = [
    np.array([[679, 361], [830, 355], [873, 573], [647, 563]], 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", imgsz=[240, 240] )
        annotated_frame = results[0].plot()
        
        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
            detections = results[0].boxes
            print('class_labels........................', class_labels)
            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}')

                            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}", (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, (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow('Annotated Frame', annotated_frame)

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

cap.release()
cv2.destroyAllWindows()