annading commited on
Commit
db1d501
·
1 Parent(s): 30130d7

changed visualizations

Browse files
Files changed (1) hide show
  1. utils.py +11 -3
utils.py CHANGED
@@ -1,10 +1,12 @@
1
  import torch
 
2
  import numpy as np
3
  import supervision as sv
4
  import cv2
5
  import os
6
  from glob import glob
7
  from tqdm import tqdm
 
8
 
9
 
10
  def plot_predictions(
@@ -25,8 +27,13 @@ def plot_predictions(
25
  in zip(labels, scores)
26
  ]
27
 
28
- bbox_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX)
29
- label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX)
 
 
 
 
 
30
  annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
31
  annotated_frame = bbox_annotator.annotate(scene=annotated_frame, detections=boxes)
32
  annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=boxes, labels=labels)
@@ -41,7 +48,8 @@ def mp4_to_png(input_path: str, save_path: str, scale_factor: float) -> str:
41
  # get frames per second
42
  fps = int(cv2.VideoCapture(input_path).get(cv2.CAP_PROP_FPS))
43
  # run subprocess to convert mp4 to pngs
44
- os.system(f"ffmpeg -i {input_path} -vf 'scale=iw*{scale_factor}:ih*{scale_factor}, fps={fps}' {save_path}/frame%08d.png")
 
45
  return fps
46
 
47
  def vid_stitcher(frames_dir: str, output_path: str, fps: int = 30) -> str:
 
1
  import torch
2
+ import subprocess
3
  import numpy as np
4
  import supervision as sv
5
  import cv2
6
  import os
7
  from glob import glob
8
  from tqdm import tqdm
9
+ import math
10
 
11
 
12
  def plot_predictions(
 
27
  in zip(labels, scores)
28
  ]
29
 
30
+ height, width, _ = image_source.shape
31
+ thickness = math.ceil(width/200)
32
+ text_scale = width/1500
33
+ text_thickness = math.ceil(text_scale*1.5)
34
+
35
+ bbox_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX, thickness=thickness)
36
+ label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX, text_scale=text_scale, text_thickness=text_thickness)
37
  annotated_frame = cv2.cvtColor(image_source, cv2.COLOR_RGB2BGR)
38
  annotated_frame = bbox_annotator.annotate(scene=annotated_frame, detections=boxes)
39
  annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=boxes, labels=labels)
 
48
  # get frames per second
49
  fps = int(cv2.VideoCapture(input_path).get(cv2.CAP_PROP_FPS))
50
  # run subprocess to convert mp4 to pngs
51
+ os.system(f"ffmpeg -i {input_path} -vf 'fps={fps},scale=iw*{scale_factor}:ih*{scale_factor}' {save_path}/frame%08d.png")
52
+ # subprocess.run(["ffmpeg", "-i", input_path, "-vf", f"scale=iw*{scale_factor}:ih*{scale_factor}, fps={fps}", f"{save_path}/frame%08d.png"])
53
  return fps
54
 
55
  def vid_stitcher(frames_dir: str, output_path: str, fps: int = 30) -> str: