| import cv2
|
| import numpy as np
|
| from ultralytics import YOLO
|
| import torch
|
|
|
| class PedestrianDetector:
|
| def __init__(self, model_path='yolov8m.pt', device=None):
|
| """
|
| Initializes the YOLOv8 detector.
|
| :param model_path: Path to the YOLOv8 model file.
|
| :param device: Device to run the model on ('cpu', 'cuda', etc.)
|
| """
|
| if device is None:
|
| self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
| else:
|
| self.device = device
|
|
|
| print(f"Initializing YOLOv8 detector on {self.device}...")
|
| self.model = YOLO(model_path)
|
| self.model.to(self.device)
|
|
|
|
|
| self.target_class = 0
|
|
|
| def detect(self, frame, conf=0.25):
|
| """
|
| Detects pedestrians in a frame.
|
| :param frame: The input image/frame.
|
| :param conf: Confidence threshold.
|
| :return: Detection results for the person class.
|
| """
|
| results = self.model.predict(
|
| source=frame,
|
| conf=conf,
|
| classes=[self.target_class],
|
| verbose=False,
|
| device=self.device
|
| )
|
| return results[0]
|
|
|
| if __name__ == "__main__":
|
|
|
| detector = PedestrianDetector()
|
| print("Detector ready.")
|
|
|
