Diego Fernandez commited on
Commit
088dea4
1 Parent(s): 3d6c603

feat: add YoloV7 Tiny support

Browse files
Files changed (3) hide show
  1. app.py +2 -2
  2. inference.py +5 -19
  3. 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=["Boxes", "Centroid"], value="Boxes", label="Detections style"
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 = 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:
@@ -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
- class ModelsPath(Enum):
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: