Videofile output for VideoAnalyser.py

Author: alexfrostdesu

DeepLabStream.py

@@ -288,7 +288,6 @@ def input_frames_for_analysis(self, frames: tuple, index: int):
                     # passes color frame to analysis
                     frame = c_frames[camera]
                     frame_time = time.time()
-                    print("input", index)
                     self._multiprocessing[camera]['input'].put((index, frame))
                     if d_maps:
                         self.store_frames(camera, frame, d_maps[camera], frame_time)
@@ -313,7 +312,6 @@ def get_analysed_frames(self) -> tuple:
                         self._start_time = time.time()  # getting the first frame here
                     # Getting the analysed data
                     analysed_index, peaks = self._multiprocessing[camera]['output'].get()
-                    print("output", analysed_index)
                     skeletons = calculate_skeletons(peaks, ANIMALS_NUMBER)
                     print('', end='\r', flush=True)  # this is the line you should not remove
 

utils/VideoAnalyser.py

@@ -48,6 +48,10 @@ def create_dataframes(data_output):
 resolution = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
 animals_list = ["Animal{}".format(num) for num in range(1, ANIMALS_NUMBER + 1)]
 
+fourcc = cv2.VideoWriter_fourcc(*'DIVX')
+output_file = os.path.join(OUT_DIR, "Analysed_" + video_name + ".avi")
+video_file = cv2.VideoWriter(output_file, fourcc, 30, resolution)
+
 
 def start_videoanalyser():
     print("Starting DeepLabCut")
@@ -77,6 +81,7 @@ def start_videoanalyser():
             else:
                 out_frame = frame
             cv2.imshow('stream', out_frame)
+            video_file.write(out_frame)
             all_rows.append(create_row(index, skeletons, experiment_enabled, experiment.get_trial()))
             index += 1
         else:
@@ -86,6 +91,8 @@ def start_videoanalyser():
             break
 
     experiment.stop_experiment()
+    video_file.release()
+    video.release()
     create_dataframes(all_rows)