Spaces:
Runtime error
Runtime error
Diego Fernandez
commited on
Commit
•
088dea4
1
Parent(s):
3d6c603
feat: add YoloV7 Tiny support
Browse files- app.py +2 -2
- inference.py +5 -19
- inference_utils.py +2 -6
app.py
CHANGED
@@ -6,14 +6,14 @@ from inference import inference
|
|
6 |
|
7 |
input_video = gr.Video(mirror_webcam=False)
|
8 |
|
9 |
-
dd_model = gr.Dropdown(choices=["YoloV7"], value="YoloV7", label="Model")
|
10 |
|
11 |
cb_motion_estimation = gr.Checkbox(value=True, label="Track camera movement")
|
12 |
|
13 |
cb_path_draw = gr.Checkbox(value=True, label="Draw objects paths")
|
14 |
|
15 |
dd_track_points = gr.Dropdown(
|
16 |
-
choices=["
|
17 |
)
|
18 |
|
19 |
slide_threshold = gr.Slider(minimum=0, maximum=1, value=0.25, label="Model confidence threshold")
|
|
|
6 |
|
7 |
input_video = gr.Video(mirror_webcam=False)
|
8 |
|
9 |
+
dd_model = gr.Dropdown(choices=["YoloV7", "YoloV7 Tiny"], value="YoloV7", label="Model")
|
10 |
|
11 |
cb_motion_estimation = gr.Checkbox(value=True, label="Track camera movement")
|
12 |
|
13 |
cb_path_draw = gr.Checkbox(value=True, label="Draw objects paths")
|
14 |
|
15 |
dd_track_points = gr.Dropdown(
|
16 |
+
choices=["Bounding box", "Centroid"], value="Bounding box", label="Detections style"
|
17 |
)
|
18 |
|
19 |
slide_threshold = gr.Slider(minimum=0, maximum=1, value=0.25, label="Model confidence threshold")
|
inference.py
CHANGED
@@ -8,13 +8,13 @@ from norfair.camera_motion import HomographyTransformationGetter, MotionEstimato
|
|
8 |
|
9 |
from inference_utils import (
|
10 |
YOLO,
|
11 |
-
ModelsPath,
|
12 |
-
Style,
|
13 |
center,
|
14 |
clean_videos,
|
15 |
draw,
|
16 |
euclidean_distance,
|
17 |
iou,
|
|
|
|
|
18 |
yolo_detections_to_norfair_detections,
|
19 |
)
|
20 |
|
@@ -22,16 +22,6 @@ DISTANCE_THRESHOLD_BBOX: float = 3.33
|
|
22 |
DISTANCE_THRESHOLD_CENTROID: int = 30
|
23 |
MAX_DISTANCE: int = 10000
|
24 |
|
25 |
-
parser = argparse.ArgumentParser(description="Track objects in a video.")
|
26 |
-
parser.add_argument("--img-size", type=int, default="720", help="YOLOv7 inference size (pixels)")
|
27 |
-
parser.add_argument(
|
28 |
-
"--iou-threshold", type=float, default="0.45", help="YOLOv7 IOU threshold for NMS"
|
29 |
-
)
|
30 |
-
parser.add_argument(
|
31 |
-
"--classes", nargs="+", type=int, help="Filter by class: --classes 0, or --classes 0 2 3"
|
32 |
-
)
|
33 |
-
args = parser.parse_args()
|
34 |
-
|
35 |
|
36 |
def inference(
|
37 |
input_video: str,
|
@@ -48,8 +38,8 @@ def inference(
|
|
48 |
coord_transformations = None
|
49 |
paths_drawer = None
|
50 |
fix_paths = False
|
51 |
-
track_points =
|
52 |
-
model = YOLO(
|
53 |
video = Video(input_path=input_video, output_path=output_path)
|
54 |
|
55 |
if motion_estimation and drawing_paths:
|
@@ -79,11 +69,7 @@ def inference(
|
|
79 |
|
80 |
for frame in video:
|
81 |
yolo_detections = model(
|
82 |
-
frame,
|
83 |
-
conf_threshold=model_threshold,
|
84 |
-
iou_threshold=args.iou_threshold,
|
85 |
-
image_size=720,
|
86 |
-
classes=args.classes,
|
87 |
)
|
88 |
|
89 |
mask = np.ones(frame.shape[:2], frame.dtype)
|
|
|
8 |
|
9 |
from inference_utils import (
|
10 |
YOLO,
|
|
|
|
|
11 |
center,
|
12 |
clean_videos,
|
13 |
draw,
|
14 |
euclidean_distance,
|
15 |
iou,
|
16 |
+
models_path,
|
17 |
+
style,
|
18 |
yolo_detections_to_norfair_detections,
|
19 |
)
|
20 |
|
|
|
22 |
DISTANCE_THRESHOLD_CENTROID: int = 30
|
23 |
MAX_DISTANCE: int = 10000
|
24 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
25 |
|
26 |
def inference(
|
27 |
input_video: str,
|
|
|
38 |
coord_transformations = None
|
39 |
paths_drawer = None
|
40 |
fix_paths = False
|
41 |
+
track_points = style[track_points]
|
42 |
+
model = YOLO(models_path[model])
|
43 |
video = Video(input_path=input_video, output_path=output_path)
|
44 |
|
45 |
if motion_estimation and drawing_paths:
|
|
|
69 |
|
70 |
for frame in video:
|
71 |
yolo_detections = model(
|
72 |
+
frame, conf_threshold=model_threshold, iou_threshold=0.45, image_size=720
|
|
|
|
|
|
|
|
|
73 |
)
|
74 |
|
75 |
mask = np.ones(frame.shape[:2], frame.dtype)
|
inference_utils.py
CHANGED
@@ -15,13 +15,9 @@ DISTANCE_THRESHOLD_CENTROID: int = 30
|
|
15 |
MAX_DISTANCE: int = 10000
|
16 |
|
17 |
|
18 |
-
|
19 |
-
YoloV7 = "models/yolov7.pt"
|
20 |
|
21 |
-
|
22 |
-
class Style(Enum):
|
23 |
-
Boxes = "bbox"
|
24 |
-
Centroid = "centroid"
|
25 |
|
26 |
|
27 |
class YOLO:
|
|
|
15 |
MAX_DISTANCE: int = 10000
|
16 |
|
17 |
|
18 |
+
models_path = {"YoloV7": "models/yolov7.pt", "YoloV7 Tiny": "models/yolov7-tiny.pt"}
|
|
|
19 |
|
20 |
+
style = {"Bounding box": "bbox", "Centroid": "centroid"}
|
|
|
|
|
|
|
21 |
|
22 |
|
23 |
class YOLO:
|