Spaces:
Runtime error
Runtime error
File size: 5,545 Bytes
fca2efd |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import argparse
import glob
import os
from enum import Enum
from typing import List, Optional, Union
import numpy as np
import torch
import torchvision.ops.boxes as bops
from norfair import norfair
from norfair.norfair import Detection
DISTANCE_THRESHOLD_BBOX: float = 3.33
DISTANCE_THRESHOLD_CENTROID: int = 30
MAX_DISTANCE: int = 10000
class ModelsPath(Enum):
YoloV7 = "models/yolov7.pt"
class Style(Enum):
Boxes = "bbox"
Centroid = "centroid"
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 euclidean_distance(detection, tracked_object):
return np.linalg.norm(detection.points - tracked_object.estimate)
def center(points):
return [np.mean(np.array(points), axis=0)]
def iou_pytorch(detection, tracked_object):
# Slower but simplier version of iou
detection_points = np.concatenate([detection.points[0], detection.points[1]])
tracked_object_points = np.concatenate([tracked_object.estimate[0], tracked_object.estimate[1]])
box_a = torch.tensor([detection_points], dtype=torch.float)
box_b = torch.tensor([tracked_object_points], dtype=torch.float)
iou = bops.box_iou(box_a, box_b)
# Since 0 <= IoU <= 1, we define 1/IoU as a distance.
# Distance values will be in [1, inf)
return np.float(1 / iou if iou else MAX_DISTANCE)
def iou(detection, tracked_object):
# Detection points will be box A
# Tracked objects point will be box B.
box_a = np.concatenate([detection.points[0], detection.points[1]])
box_b = np.concatenate([tracked_object.estimate[0], tracked_object.estimate[1]])
x_a = max(box_a[0], box_b[0])
y_a = max(box_a[1], box_b[1])
x_b = min(box_a[2], box_b[2])
y_b = min(box_a[3], box_b[3])
# Compute the area of intersection rectangle
inter_area = max(0, x_b - x_a + 1) * max(0, y_b - y_a + 1)
# Compute the area of both the prediction and tracker
# rectangles
box_a_area = (box_a[2] - box_a[0] + 1) * (box_a[3] - box_a[1] + 1)
box_b_area = (box_b[2] - box_b[0] + 1) * (box_b[3] - box_b[1] + 1)
# Compute the intersection over union by taking the intersection
# area and dividing it by the sum of prediction + tracker
# areas - the interesection area
iou = inter_area / float(box_a_area + box_b_area - inter_area)
# Since 0 <= IoU <= 1, we define 1/IoU as a distance.
# Distance values will be in [1, inf)
return 1 / iou if iou else (MAX_DISTANCE)
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()])
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
def clean_videos(path: str):
# Remove past videos
files = glob.glob(f"{path}/*")
for file in files:
if file.endswith(".mp4"):
os.remove(file)
def draw(paths_drawer, track_points, frame, detections, tracked_objects):
if track_points == "centroid":
norfair.draw_points(frame, detections)
norfair.draw_tracked_objects(frame, tracked_objects)
elif track_points == "bbox":
norfair.draw_boxes(frame, detections)
norfair.draw_tracked_boxes(frame, tracked_objects)
if paths_drawer is not None:
frame = paths_drawer.draw(frame, tracked_objects)
return frame
|