Spaces:
Runtime error
Runtime error
File size: 2,953 Bytes
fca2efd 00738ec fca2efd 210ae8c 3647577 fca2efd 210ae8c 088dea4 fca2efd 210ae8c fca2efd 210ae8c fca2efd 00738ec fca2efd 9d1a8a7 088dea4 367d735 fca2efd 210ae8c 9d1a8a7 fca2efd 30935b4 fca2efd 210ae8c fca2efd 9d1a8a7 fca2efd 088dea4 fca2efd 9d1a8a7 fca2efd 367d735 |
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 |
import argparse
import glob
import os
import tempfile
import numpy as np
from norfair import AbsolutePaths, Paths, Tracker, Video
from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
from custom_models import YOLO, yolo_detections_to_norfair_detections
from demo_utils.configuration import (
DISTANCE_THRESHOLD_BBOX,
DISTANCE_THRESHOLD_CENTROID,
models_path,
style,
)
from demo_utils.distance_function import euclidean_distance, iou
from demo_utils.draw import center, draw
def inference(
input_video: str,
model: str,
features: str,
track_points: str,
model_threshold: str,
):
temp_dir = tempfile.TemporaryDirectory()
output_path = temp_dir.name
coord_transformations = None
paths_drawer = None
fix_paths = False
track_points = style[track_points]
model = YOLO(models_path[model])
video = Video(input_path=input_video, output_path=output_path)
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 and drawing_paths:
fix_paths = True
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 == "bbox" else euclidean_distance
distance_threshold = (
DISTANCE_THRESHOLD_BBOX if track_points == "bbox" else DISTANCE_THRESHOLD_CENTROID
)
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=5, thickness=2)
for frame in video:
yolo_detections = model(
frame, conf_threshold=model_threshold, iou_threshold=0.45, image_size=720
)
mask = np.ones(frame.shape[:2], frame.dtype)
if motion_estimation:
coord_transformations = motion_estimator.update(frame, mask)
detections = yolo_detections_to_norfair_detections(
yolo_detections, track_points=track_points
)
tracked_objects = tracker.update(
detections=detections, coord_transformations=coord_transformations
)
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 os.path.join(output_path, file_name)
|