(path, video, fps=25, scale=2, audio=None, audio_rate=16000, overlay_pts=None, ffmpeg_experimental=False)
| 83 | |
| 84 | |
| 85 | def save_video(path, video, fps=25, scale=2, audio=None, audio_rate=16000, overlay_pts=None, ffmpeg_experimental=False): |
| 86 | if not os.path.exists(os.path.dirname(path)): |
| 87 | os.makedirs(os.path.dirname(path)) |
| 88 | success = True |
| 89 | out_size = (scale * video.shape[-1], scale * video.shape[-2]) |
| 90 | video_path = get_temp_path(os.path.split(path)[0], ext=".mp4") |
| 91 | if torch.is_tensor(video): |
| 92 | vid = video.squeeze().detach().cpu().numpy() |
| 93 | else: |
| 94 | vid = video.copy() # Make a copy so that we don't alter the object |
| 95 | |
| 96 | if np.min(vid) < 0: |
| 97 | vid = 127 * vid + 127 |
| 98 | elif np.max(vid) <= 1: |
| 99 | vid = 255 * vid |
| 100 | |
| 101 | is_color = True |
| 102 | if vid.ndim == 3: |
| 103 | is_color = False |
| 104 | |
| 105 | writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), float(fps), out_size, isColor=is_color) |
| 106 | for i, frame in enumerate(vid): |
| 107 | if is_color: |
| 108 | frame = cv2.cvtColor(np.rollaxis(frame, 0, 3), cv2.COLOR_RGB2BGR) |
| 109 | |
| 110 | if scale != 1: |
| 111 | frame = cv2.resize(frame, out_size) |
| 112 | |
| 113 | write_frame = frame.astype('uint8') |
| 114 | |
| 115 | if overlay_pts is not None: |
| 116 | for pt in overlay_pts[i]: |
| 117 | cv2.circle(write_frame, (int(scale * pt[0]), int(scale * pt[1])), 2, (0, 0, 0), -1) |
| 118 | |
| 119 | writer.write(write_frame) |
| 120 | writer.release() |
| 121 | |
| 122 | inputs = [ffmpeg.input(video_path)['v']] |
| 123 | |
| 124 | if audio is not None: # Save the audio file |
| 125 | audio_path = swp_extension(video_path, ".wav") |
| 126 | save_audio(audio_path, audio, audio_rate) |
| 127 | inputs += [ffmpeg.input(audio_path)['a']] |
| 128 | |
| 129 | try: |
| 130 | if ffmpeg_experimental: |
| 131 | out = ffmpeg.output(*inputs, path, strict='-2', loglevel="panic", vcodec='h264').overwrite_output() |
| 132 | else: |
| 133 | out = ffmpeg.output(*inputs, path, loglevel="panic", vcodec='h264').overwrite_output() |
| 134 | out.run(quiet=True) |
| 135 | except: |
| 136 | success = False |
| 137 | |
| 138 | if audio is not None and os.path.isfile(audio_path): |
| 139 | os.remove(audio_path) |
| 140 | if os.path.isfile(video_path): |
| 141 | os.remove(video_path) |
| 142 |
no test coverage detected