|
import os |
|
from typing import List, Optional, Union |
|
|
|
import numpy as np |
|
import torch |
|
from norfair import Detection |
|
|
|
|
|
class YOLO: |
|
def __init__(self, model_path: str, device: Optional[str] = None): |
|
if device is not None and "cuda" in device and not torch.cuda.is_available(): |
|
raise Exception("Selected device='cuda', but cuda is not available to Pytorch.") |
|
|
|
elif device is None: |
|
device = "cuda:0" if torch.cuda.is_available() else "cpu" |
|
|
|
if not os.path.exists(model_path): |
|
os.system( |
|
f"wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/{os.path.basename(model_path)} -O {model_path}" |
|
) |
|
|
|
|
|
try: |
|
self.model = torch.hub.load("WongKinYiu/yolov7", "custom", model_path) |
|
except: |
|
raise Exception("Failed to load model from {}".format(model_path)) |
|
|
|
def __call__( |
|
self, |
|
img: Union[str, np.ndarray], |
|
conf_threshold: float = 0.25, |
|
iou_threshold: float = 0.45, |
|
image_size: int = 720, |
|
classes: Optional[List[int]] = None, |
|
) -> torch.tensor: |
|
|
|
self.model.conf = conf_threshold |
|
self.model.iou = iou_threshold |
|
if classes is not None: |
|
self.model.classes = classes |
|
detections = self.model(img, size=image_size) |
|
return detections |
|
|
|
|
|
def yolo_detections_to_norfair_detections( |
|
yolo_detections: torch.tensor, track_points: str = "centroid" |
|
) -> List[Detection]: |
|
"""convert detections_as_xywh to norfair detections""" |
|
norfair_detections: List[Detection] = [] |
|
|
|
if track_points == "centroid": |
|
detections_as_xywh = yolo_detections.xywh[0] |
|
for detection_as_xywh in detections_as_xywh: |
|
centroid = np.array([detection_as_xywh[0].item(), detection_as_xywh[1].item()]) |
|
scores = np.array([detection_as_xywh[4].item()]) |
|
norfair_detections.append(Detection(points=centroid, scores=scores)) |
|
elif track_points == "bbox": |
|
detections_as_xyxy = yolo_detections.xyxy[0] |
|
for detection_as_xyxy in detections_as_xyxy: |
|
bbox = np.array( |
|
[ |
|
[detection_as_xyxy[0].item(), detection_as_xyxy[1].item()], |
|
[detection_as_xyxy[2].item(), detection_as_xyxy[3].item()], |
|
] |
|
) |
|
scores = np.array([detection_as_xyxy[4].item(), detection_as_xyxy[4].item()]) |
|
norfair_detections.append(Detection(points=bbox, scores=scores)) |
|
|
|
return norfair_detections |
|
|