from ultralytics import YOLO
import torch

class PedestrianTracker:
    def __init__(self, model_path='yolov8m.pt', device=None):
        if device is None:
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device
            
        self.model = YOLO(model_path)
        self.model.to(self.device)

    def track(self, frame, conf=0.25, persist=True):
        """
        Tracks multiple classes: People (0), Bicycle (1), Car (2), Motorcycle (3), Bus (5), Truck (7).
        """
        results = self.model.track(
            source=frame,
            conf=conf,
            persist=persist,
            tracker="bytetrack.yaml",
            classes=[0, 1, 2, 3, 5, 7], # Person + All major vehicles
            verbose=False,
            device=self.device
        )
        return results[0]