norfair-demo / inference.py
Diego Fernández
chore: change distance function thresholds
de42222
raw
history blame
3.55 kB
import numpy as np
from norfair import AbsolutePaths, Paths, Tracker, Video
from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
from norfair.distances import create_normalized_mean_euclidean_distance
from custom_models import YOLO, yolo_detections_to_norfair_detections
from demo_utils.configuration import (
DISTANCE_THRESHOLD_BBOX,
DISTANCE_THRESHOLD_CENTROID,
examples,
models_path,
style,
)
from demo_utils.draw import center, draw
def inference(
input_video: str,
model: str = "YOLOv7",
features: str = [0, 1],
track_points: str = "Bounding box",
model_threshold: float = 0.25,
):
coord_transformations = None
paths_drawer = None
fix_paths = False
classes = None
track_points = style[track_points]
model = YOLO(models_path[model])
video = Video(input_path=input_video)
motion_estimation = len(features) > 0 and (
features[0] == 0 or (len(features) > 1 and features[1] == 0)
)
drawing_paths = len(features) > 0 and (
features[0] == 1 or (len(features) > 1 and features[1] == 1)
)
if motion_estimation:
transformations_getter = HomographyTransformationGetter()
motion_estimator = MotionEstimator(
max_points=500, min_distance=7, transformations_getter=transformations_getter
)
distance_function = "iou" if track_points == style["Bounding box"] else "euclidean"
distance_threshold = (
DISTANCE_THRESHOLD_BBOX
if track_points == style["Bounding box"]
else DISTANCE_THRESHOLD_CENTROID
)
if motion_estimation and drawing_paths:
fix_paths = True
# Examples configuration
for example in examples:
if example not in input_video:
continue
fix_paths = examples[example]["absolute_path"]
distance_threshold = examples[example]["distance_threshold"]
classes = examples[example]["classes"]
print(f"Set config to {example}: {fix_paths} {distance_threshold} {classes}")
break
tracker = Tracker(
distance_function=distance_function,
distance_threshold=distance_threshold,
)
if drawing_paths:
paths_drawer = Paths(center, attenuation=0.01)
if fix_paths:
paths_drawer = AbsolutePaths(max_history=50, thickness=2)
for frame in video:
yolo_detections = model(
frame,
conf_threshold=model_threshold,
iou_threshold=0.45,
image_size=720,
classes=classes,
)
detections = yolo_detections_to_norfair_detections(
yolo_detections, track_points=track_points
)
tracked_objects = tracker.update(
detections=detections, coord_transformations=coord_transformations
)
if motion_estimation:
mask = np.ones(frame.shape[:2], frame.dtype)
if track_points == "bbox":
for det in detections:
i = det.points.astype(int)
mask[i[0, 1] : i[1, 1], i[0, 0] : i[1, 0]] = 0
coord_transformations = motion_estimator.update(frame, mask)
frame = draw(
paths_drawer,
track_points,
frame,
detections,
tracked_objects,
coord_transformations,
fix_paths,
)
video.write(frame)
base_file_name = input_video.split("/")[-1].split(".")[0]
file_name = base_file_name + "_out.mp4"
return file_name