Spaces:
Runtime error
Runtime error
File size: 2,816 Bytes
210ae8c 57a060c 210ae8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
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.")
# automatically set device if its None
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}"
)
# load model
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" # bbox or 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()],
[detection_as_xywh[0].item(), detection_as_xywh[1].item()],
]
)
scores = np.array([detection_as_xywh[4].item(), 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
|