Spaces:
Runtime error
Runtime error
Diego Fernandez
commited on
Commit
·
9d1a8a7
1
Parent(s):
dfb84ad
feat: path with moving camera
Browse files- inference.py +17 -2
- inference_utils.py +13 -3
inference.py
CHANGED
@@ -3,7 +3,7 @@ import glob
|
|
3 |
import os
|
4 |
|
5 |
import numpy as np
|
6 |
-
from norfair import Paths, Tracker, Video
|
7 |
from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
|
8 |
|
9 |
from inference_utils import (
|
@@ -47,10 +47,14 @@ def inference(
|
|
47 |
|
48 |
coord_transformations = None
|
49 |
paths_drawer = None
|
|
|
50 |
track_points = Style[track_points].value
|
51 |
model = YOLO(ModelsPath[model].value)
|
52 |
video = Video(input_path=input_video, output_path=output_path)
|
53 |
|
|
|
|
|
|
|
54 |
if motion_estimation:
|
55 |
transformations_getter = HomographyTransformationGetter()
|
56 |
|
@@ -70,6 +74,9 @@ def inference(
|
|
70 |
if drawing_paths:
|
71 |
paths_drawer = Paths(center, attenuation=0.01)
|
72 |
|
|
|
|
|
|
|
73 |
for frame in video:
|
74 |
yolo_detections = model(
|
75 |
frame,
|
@@ -92,7 +99,15 @@ def inference(
|
|
92 |
detections=detections, coord_transformations=coord_transformations
|
93 |
)
|
94 |
|
95 |
-
frame = draw(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
96 |
video.write(frame)
|
97 |
|
98 |
base_file_name = input_video.split("/")[-1].split(".")[0]
|
|
|
3 |
import os
|
4 |
|
5 |
import numpy as np
|
6 |
+
from norfair import AbsolutePaths, FixedCamera, Paths, Tracker, Video
|
7 |
from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
|
8 |
|
9 |
from inference_utils import (
|
|
|
47 |
|
48 |
coord_transformations = None
|
49 |
paths_drawer = None
|
50 |
+
fix_paths = False
|
51 |
track_points = Style[track_points].value
|
52 |
model = YOLO(ModelsPath[model].value)
|
53 |
video = Video(input_path=input_video, output_path=output_path)
|
54 |
|
55 |
+
if motion_estimation and drawing_paths:
|
56 |
+
fix_paths = True
|
57 |
+
|
58 |
if motion_estimation:
|
59 |
transformations_getter = HomographyTransformationGetter()
|
60 |
|
|
|
74 |
if drawing_paths:
|
75 |
paths_drawer = Paths(center, attenuation=0.01)
|
76 |
|
77 |
+
if fix_paths:
|
78 |
+
paths_drawer = AbsolutePaths(max_history=5, thickness=2)
|
79 |
+
|
80 |
for frame in video:
|
81 |
yolo_detections = model(
|
82 |
frame,
|
|
|
99 |
detections=detections, coord_transformations=coord_transformations
|
100 |
)
|
101 |
|
102 |
+
frame = draw(
|
103 |
+
paths_drawer,
|
104 |
+
track_points,
|
105 |
+
frame,
|
106 |
+
detections,
|
107 |
+
tracked_objects,
|
108 |
+
coord_transformations,
|
109 |
+
fix_paths,
|
110 |
+
)
|
111 |
video.write(frame)
|
112 |
|
113 |
base_file_name = input_video.split("/")[-1].split(".")[0]
|
inference_utils.py
CHANGED
@@ -8,7 +8,7 @@ import norfair
|
|
8 |
import numpy as np
|
9 |
import torch
|
10 |
import torchvision.ops.boxes as bops
|
11 |
-
from norfair import Detection
|
12 |
|
13 |
DISTANCE_THRESHOLD_BBOX: float = 3.33
|
14 |
DISTANCE_THRESHOLD_CENTROID: int = 30
|
@@ -148,7 +148,15 @@ def clean_videos(path: str):
|
|
148 |
os.remove(file)
|
149 |
|
150 |
|
151 |
-
def draw(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
152 |
if track_points == "centroid":
|
153 |
norfair.draw_points(frame, detections)
|
154 |
norfair.draw_tracked_objects(frame, tracked_objects)
|
@@ -156,7 +164,9 @@ def draw(paths_drawer, track_points, frame, detections, tracked_objects):
|
|
156 |
norfair.draw_boxes(frame, detections)
|
157 |
norfair.draw_tracked_boxes(frame, tracked_objects)
|
158 |
|
159 |
-
if
|
|
|
|
|
160 |
frame = paths_drawer.draw(frame, tracked_objects)
|
161 |
|
162 |
return frame
|
|
|
8 |
import numpy as np
|
9 |
import torch
|
10 |
import torchvision.ops.boxes as bops
|
11 |
+
from norfair import Detection, draw_absolute_grid
|
12 |
|
13 |
DISTANCE_THRESHOLD_BBOX: float = 3.33
|
14 |
DISTANCE_THRESHOLD_CENTROID: int = 30
|
|
|
148 |
os.remove(file)
|
149 |
|
150 |
|
151 |
+
def draw(
|
152 |
+
paths_drawer,
|
153 |
+
track_points,
|
154 |
+
frame,
|
155 |
+
detections,
|
156 |
+
tracked_objects,
|
157 |
+
coord_transformations,
|
158 |
+
fix_paths,
|
159 |
+
):
|
160 |
if track_points == "centroid":
|
161 |
norfair.draw_points(frame, detections)
|
162 |
norfair.draw_tracked_objects(frame, tracked_objects)
|
|
|
164 |
norfair.draw_boxes(frame, detections)
|
165 |
norfair.draw_tracked_boxes(frame, tracked_objects)
|
166 |
|
167 |
+
if fix_paths:
|
168 |
+
frame = paths_drawer.draw(frame, tracked_objects, coord_transformations)
|
169 |
+
elif paths_drawer is not None:
|
170 |
frame = paths_drawer.draw(frame, tracked_objects)
|
171 |
|
172 |
return frame
|