improve video generation and logging tqdm

app.py

@@ -8,6 +8,7 @@ from typing import Dict, List, Tuple
 
 import datetime
 import logging
+import warnings
 
 # imports for tracking
 import torch
@@ -70,11 +71,14 @@ config_track = DotDict({
 
 
 logger.info('---Yolo model...')
+# Yolo has warning problems, so we set an env variable to remove it
+os.environ["VERBOSE"] = "False"
 URL_MODEL = "https://github.com/surfriderfoundationeurope/IA_Pau/releases/download/v0.1/yolov5.pt"
 FILE_MODEL = "yolov5.pt"
 model_path = download_from_url(URL_MODEL, FILE_MODEL, logger, "./models")
 model_yolo = load_model(model_path, config_track.device)
 
+
 logger.info('---Centernet model...')
 URL_MODEL = "https://partage.imt.fr/index.php/s/sJi22N6gedN6T4q/download"
 FILE_MODEL = "mobilenet_v3_pretrained.pth"
@@ -132,8 +136,11 @@ def track(args):
     detections = []
     logger.info('---Detecting...')
     if args.model_type == "yolo":
-        for frame in reader:
-            detections.append(detector(frame))
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore")
+
+            for frame in reader:
+                detections.append(detector(frame))
     elif args.model_type == "centernet":
         detections = get_detections_for_video(reader, detector, batch_size=args.detection_batch_size, device=device)
 
@@ -141,7 +148,7 @@ def track(args):
     display = None
     results = track_video(reader, iter(detections), args, engine, transition_variance, observation_variance, display, is_yolo=args.model_type=="yolo")
     reader.video.release()
-    
+
     # store unfiltered results
     datestr = datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
     output_filename = op.splitext(args.video_path)[0] + "_" + datestr + '_unfiltered.txt'
@@ -179,11 +186,15 @@ def run_model(video_path, model_type, seconds, skip, tau, kappa):
     with open(output_json_path, 'w') as f_out:
         json.dump(output_json, f_out)
 
-    #output_json={"detected_trash":[]}
     # build video output
     logger.info('---Generating new video...')
-    video = SimpleVideoReader(video_path, skip_frames=0)
-    generate_video_with_annotations(video, output_json, output_path,
+    reader = IterableFrameReader(video_filename=config_track.video_path,
+                                 skip_frames=0,
+                                 progress_bar=True,
+                                 preload=False,
+                                 max_frame=config_track.max_length)
+
+    generate_video_with_annotations(reader, output_json, output_path,
                                     config_track.skip_frames, config_track.max_length,
                                     config_track.downscale_output, logger)
 

requirements.txt

@@ -1,4 +1,5 @@
 scikit-image
+scikit-video
 opencv-python
 torch
 torchvision
@@ -9,7 +10,6 @@ tqdm
 tensorboard
 imgaug
 psycopg2-binary
-moviepy
 git+https://github.com/pykalman/pykalman.git
 yolov5
 gradio

tools/video_readers.py

@@ -3,6 +3,7 @@ import torch
 from tqdm import tqdm
 from itertools import cycle
 
+
 class AdvancedFrameReader:
     def __init__(self, video_name, read_every, rescale_factor, init_time_min, init_time_s):
 
@@ -18,7 +19,6 @@ class AdvancedFrameReader:
 
         self.init_rescale_factor = rescale_factor
 
-
         self.frame_skip = read_every -  1
         self.fps = self.cap.get(cv2.CAP_PROP_FPS)/read_every
         print(f'Reading at {self.fps:.2f} fps')
@@ -27,7 +27,6 @@ class AdvancedFrameReader:
         self.init_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
         self.total_num_frames = self.cap.get(cv2.CAP_PROP_FRAME_COUNT)
 
-
     def post_process(self, ret, frame):
         if ret:
             if self.original_shape_mode:
@@ -55,7 +54,7 @@ class AdvancedFrameReader:
         time = 60 * time_min  + time_s
         self.cap.set(cv2.CAP_PROP_POS_MSEC, 1000 * time)
         self.init_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
-        print('Reading from {}min{}sec'.format(time_min,time_s))
+        #print('Reading from {}min{}sec'.format(time_min,time_s))
         self.nb_frames_read = 0
 
     def reset_init_frame(self):
@@ -72,7 +71,7 @@ class AdvancedFrameReader:
         height = int(self.original_height/rescale_factor)
         self.new_shape  = (width, height)
         self.original_shape_mode = False
-        print('Reading in {}x{}'.format(width, height))
+        #print('Reading in {}x{}'.format(width, height))
 
     def set_original_shape_mode(self, mode):
         self.original_shape_mode = mode
@@ -83,17 +82,21 @@ class AdvancedFrameReader:
 
 class IterableFrameReader:
     def __init__(self, video_filename, skip_frames=0, output_shape=None, progress_bar=False, preload=False, max_frame=0):
+        # store arguments for reset
         self.video_filename = video_filename
         self.max_frame_arg = max_frame
+        self.progress_bar_arg = progress_bar
+        self.preload = preload
+        self.skip_frames = skip_frames
 
         self.video = cv2.VideoCapture(video_filename)
-        self.input_shape = (self.video.get(cv2.CAP_PROP_FRAME_WIDTH), self.video.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        self.skip_frames = skip_frames
-        self.preload = preload
-        self.total_num_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)
+        self.input_shape = (int(self.video.get(cv2.CAP_PROP_FRAME_WIDTH)),
+                            int(self.video.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+        self.total_num_frames = int(self.video.get(cv2.CAP_PROP_FRAME_COUNT))
 
         self.max_num_frames = min(max_frame, self.total_num_frames) if max_frame!=0 else self.total_num_frames
         self.counter = 0
+        self.progress_bar = None
 
         if output_shape is None:
             w, h = self.input_shape
@@ -105,24 +108,27 @@ class IterableFrameReader:
 
         self.fps = self.video.get(cv2.CAP_PROP_FPS) / (self.skip_frames+1)
 
-        print(f'Reading video at {self.fps}fps.')
-        if progress_bar:
-            self.progress_bar = tqdm(total=int(self.max_num_frames/(self.skip_frames+1)), leave=True)
-            self.progress_bar_update = self.progress_bar.update
-        else:
-            self.progress_bar_update = lambda: None
-
         if self.preload:
-            print('Preloading frames in RAM...')
             self.frames = self._load_all_frames()
 
+    def update_progress_bar(self):
+        if self.progress_bar_arg:
+            if self.progress_bar:
+                # update_progress_bar
+                self.progress_bar.update()
+            else:
+                # create progress bar
+                self.progress_bar = tqdm(total=int(self.max_num_frames/(self.skip_frames+1)),
+                                         position=1, leave=True)
+
     def reset_video(self):
         """ This method is needed as cv2.CAP_PROP_POS_FRAMES
         does not work on all backends
         """
         self.video.release()
+        self.progress_bar.close()
         self.__init__(self.video_filename, self.skip_frames, self.output_shape,
-                      self.progress_bar is not None, self.preload, self.max_frame_arg)
+                      self.progress_bar_arg, self.preload, self.max_frame_arg)
 
     def _load_all_frames(self):
         frames = []
@@ -139,7 +145,7 @@ class IterableFrameReader:
         if self.preload:
             if self.counter < len(self.frames):
                 frame = self.frames[self.counter]
-                self.progress_bar_update()
+                self.update_progress_bar()
                 return frame
         else:
             if self.counter < self.max_num_frames:
@@ -150,12 +156,11 @@ class IterableFrameReader:
         self.reset_video()
         raise StopIteration
 
-
     def _read_frame(self):
         ret, frame = self.video.read()
         self._skip_frames()
         if ret:
-            self.progress_bar_update()
+            self.update_progress_bar()
             frame =  cv2.resize(frame, self.output_shape)
         return ret, frame
 
@@ -191,8 +196,8 @@ class SimpleVideoReader:
         for _ in range(self.skip_frames):
             self.video.read()
 
-class TorchIterableFromReader(torch.utils.data.IterableDataset):
 
+class TorchIterableFromReader(torch.utils.data.IterableDataset):
     def __init__(self, reader, transforms):
         self.transforms = transforms
         self.reader = reader

tracking/utils.py

@@ -8,8 +8,8 @@ from tools.video_readers import TorchIterableFromReader
 from time import time
 from detection.transforms import TransformFrames
 from collections import defaultdict
-from moviepy.editor import ImageSequenceClip
 from skimage.transform import downscale_local_mean
+from skvideo.io import FFmpegWriter
 
 class GaussianMixture(object):
     def __init__(self, means, covariance, weights):
@@ -82,7 +82,7 @@ def get_detections_for_video(reader, detector, batch_size=16, device=None):
     return detections
 
 
-def generate_video_with_annotations(video, output_detected, output_filename, skip_frames, maxframes, downscale, logger):
+def generate_video_with_annotations(reader, output_detected, output_filename, skip_frames, maxframes, downscale, logger):
     fps = 24
     logger.info("---intepreting json")
     results = defaultdict(list)
@@ -105,31 +105,30 @@ def generate_video_with_annotations(video, output_detected, output_filename, ski
                     results[frame_nb * (skip_frames+1) + i].append((object_nb, new_x, new_y, object_class))
     logger.info("---writing video")
 
-    #fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    # writer = cv2.VideoWriter(filename=output_filename,
-                                    #apiPreference=cv2.CAP_FFMPEG,
-    #                                fourcc=fourcc,
-    #                                fps=fps,
-    #                                frameSize=video.shape)
+    writer = FFmpegWriter(filename = output_filename)
 
     font = cv2.FONT_HERSHEY_COMPLEX
-    ret, frame, frame_nb = video.read()
-    frames = []
-    while ret:
+    for frame_nb, frame in enumerate(reader):
         detections_for_frame = results[frame_nb]
         for detection in detections_for_frame:
             cv2.putText(frame, f'{detection[0]}/{detection[3]}', (int(detection[1]), int(detection[2])+5), font, 2, (0, 0, 255), 3, cv2.LINE_AA)
 
         frame = downscale_local_mean(frame, (downscale,downscale,1)).astype(np.uint8)
-        frames.append(frame[:,:,::-1])
+        writer.writeFrame(frame[:,:,::-1])
+        # moviepy version
+        # frames.append(frame[:,:,::-1])
 
-        ret, frame, frame_nb = video.read()
-        if frame_nb > maxframes:
-            break
+        #ret, frame, frame_nb = video.read()
+        #if frame_nb > maxframes:
+        #    break
 
-    clip = ImageSequenceClip(sequence=frames, fps=fps)
-    clip.write_videofile(output_filename, fps=fps)
-    del frames
+    writer.close()
+    reader.video.release()
+
+    # version with moviepy
+    #clip = ImageSequenceClip(sequence=frames, fps=fps)
+    #clip.write_videofile(output_filename, fps=fps)
+    #del frames
 
     logger.info("---finished writing video")