Diego Fernandez commited on
Commit
367d735
1 Parent(s): 1dea2eb

chore: improve returned path

Browse files
Files changed (1) hide show
  1. inference.py +7 -3
inference.py CHANGED
@@ -41,13 +41,15 @@ def inference(
41
  track_points: str,
42
  model_threshold: str,
43
  ):
44
- clean_videos("tmp")
 
 
45
 
46
  coord_transformations = None
47
  paths_drawer = None
48
  track_points = Style[track_points].value
49
  model = YOLO(ModelsPath[model].value)
50
- video = Video(input_path=input_video, output_path="tmp")
51
 
52
  if motion_estimation:
53
  transformations_getter = HomographyTransformationGetter()
@@ -96,4 +98,6 @@ def inference(
96
  frame = draw(paths_drawer, track_points, frame, detections, tracked_objects)
97
  video.write(frame)
98
 
99
- return f"{input_video[1:-4]}_out.mp4"
 
 
 
41
  track_points: str,
42
  model_threshold: str,
43
  ):
44
+ output_path = "/tmp"
45
+
46
+ clean_videos(output_path)
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()
 
98
  frame = draw(paths_drawer, track_points, frame, detections, tracked_objects)
99
  video.write(frame)
100
 
101
+ base_file_name = input_video.split("/")[-1].split(".")[0]
102
+ file_name = base_file_name + "_out.mp4"
103
+ return os.path.join(output_path, file_name)