Update utils/dc_utils.py

utils/dc_utils.py

@@ -6,11 +6,11 @@
 import numpy as np
 import matplotlib.cm as cm
 import imageio
+import cv2
 try:
     from decord import VideoReader, cpu
     DECORD_AVAILABLE = True
 except:
-    import cv2
     DECORD_AVAILABLE = False
 
 def ensure_even(value):
@@ -22,12 +22,10 @@ def read_video_frames(video_path, process_length, target_fps=-1, max_res=-1):
         original_height, original_width = vid.get_batch([0]).shape[1:3]
         height = original_height
         width = original_width
-        print(f'==> original video size: {original_height} x {original_width}')
         if max_res > 0 and max(height, width) > max_res:
             scale = max_res / max(original_height, original_width)
             height = ensure_even(round(original_height * scale))
             width = ensure_even(round(original_width * scale))
-            print(f'==> downsample video size: {height} x {width}')
 
         vid = VideoReader(video_path, ctx=cpu(0), width=width, height=height)
 
@@ -71,18 +69,13 @@ def read_video_frames(video_path, process_length, target_fps=-1, max_res=-1):
     return frames, fps
 
 
-def save_video(frames, output_video_path, fps=10, is_depths=False, grayscale=False):
+def save_video(frames, depths, output_video_path, fps=10, grayscale=False):
     writer = imageio.get_writer(output_video_path, fps=fps, macro_block_size=1, codec='libx264', ffmpeg_params=['-crf', '18'])
-    if is_depths:
-        colormap = np.array(cm.get_cmap("inferno").colors)
-        d_min, d_max = frames.min(), frames.max()
-        for i in range(frames.shape[0]):
-            depth = frames[i]
-            depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8)
-            depth_vis = (colormap[depth_norm] * 255).astype(np.uint8) if not grayscale else depth_norm
-            writer.append_data(depth_vis)
-    else:
-        for i in range(frames.shape[0]):
-            writer.append_data(frames[i])
-
+    colormap = np.array(cm.get_cmap("inferno").colors)
+    d_min, d_max = depths.min(), depths.max()
+    for i in range(frames.shape[0]):
+        depth = depths[i]
+        depth_norm = ((depth - d_min) / (d_max - d_min) * 255).astype(np.uint8)
+        depth_vis = (colormap[depth_norm] * 255).astype(np.uint8) if not grayscale else cv2.cvtColor(depth_norm, cv2.COLOR_GRAY2BGR)
+        writer.append_data(np.hstack((frames[i], depth_vis)))
     writer.close()