object_tracking_final.py

import cv2
import time
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort

# Load YOLO model
model = YOLO("yolov8n.pt")

# Initialize tracker
tracker = DeepSort(
    max_age=50,
    n_init=2,
    max_iou_distance=0.8
)

cap = cv2.VideoCapture(0)

cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)

cv2.namedWindow(
    "Final Object Tracking System",
    cv2.WINDOW_NORMAL
)

cv2.resizeWindow(
    "Final Object Tracking System",
    1200,
    800
)

# Video recording setup
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

out = cv2.VideoWriter(
    "tracked_video.mp4",
    cv2.VideoWriter_fourcc(*'mp4v'),
    20,
    (frame_width, frame_height)
)


prev_time = time.time()

allowed_classes = [
    "person",
    "cell phone",
    "bottle",
    "chair",
    "laptop"
]

cv2.namedWindow(
    "Final Object Tracking System",
    cv2.WINDOW_NORMAL
)

cv2.resizeWindow(
    "Final Object Tracking System",
    1200,
    800
)

while True:

    ret, frame = cap.read()

    if not ret:
        break

    results = model(frame)

    detections = []

    for result in results:

        for box in result.boxes:

            confidence = float(box.conf[0])

            if confidence < 0.6:
                continue

            class_id = int(box.cls[0])
            class_name = model.names[class_id]

            if class_name not in allowed_classes:
                continue

            x1, y1, x2, y2 = map(int, box.xyxy[0])

            width = x2 - x1
            height = y2 - y1

            detections.append(
                ([x1, y1, width, height], confidence, class_name)
            )

    tracks = tracker.update_tracks(
        detections,
        frame=frame
    )

    object_count = 0

    for track in tracks:

        if not track.is_confirmed():
            continue

        object_count += 1

        track_id = track.track_id

        x1, y1, x2, y2 = map(int, track.to_ltrb())

        cv2.rectangle(
            frame,
            (x1, y1),
            (x2, y2),
            (0, 255, 0),
            2
        )

        cv2.putText(
            frame,
            f"ID: {track_id}",
            (x1, y1 - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 255, 0),
            2
        )

    current_time = time.time()
    fps = 1 / (current_time - prev_time)
    prev_time = current_time

    cv2.putText(
        frame,
        f"FPS: {int(fps)}",
        (20, 40),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        (0, 255, 255),
        2
    )

    cv2.putText(
        frame,
        f"Objects: {object_count}",
        (20, 80),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        (0, 255, 255),
        2
    )

    out.write(frame)

    cv2.imshow("Final Object Tracking System", frame)

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

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

print("Video saved as tracked_video.mp4")