| 205 | |
| 206 | |
| 207 | def transcode( |
| 208 | frames, |
| 209 | input_args: list = ["pipe:"], |
| 210 | input_kwargs: dict = { |
| 211 | "format": "rawvideo", |
| 212 | "pix_fmt": "rgb24", |
| 213 | "framerate": MPEG_TRANSCODING_FRAMERATE, |
| 214 | }, |
| 215 | output_args: list = ["pipe:"], |
| 216 | output_kwargs: dict = { |
| 217 | "format": "mpeg", |
| 218 | "vcodec": "mpeg1video", |
| 219 | "qscale": 2, # Set the quality scale (lower is better quality) |
| 220 | "video_bitrate": "5000k", |
| 221 | }, |
| 222 | ) -> bytes: |
| 223 | t, h, w, c = frames.shape |
| 224 | assert c == 3, f"{c=}" |
| 225 | process = ffmpeg.input( |
| 226 | *input_args, |
| 227 | **input_kwargs, |
| 228 | s="{}x{}".format(w, h), |
| 229 | ) |
| 230 | process = process.output( |
| 231 | *output_args, |
| 232 | **output_kwargs, |
| 233 | ) |
| 234 | process = process.overwrite_output() |
| 235 | process = process.run_async(pipe_stdin=True, pipe_stdout=True, pipe_stderr=True) |
| 236 | stdout_data, stderr_data = process.communicate(input=frames.tobytes()) |
| 237 | stderr = stderr_data.decode("utf-8") |
| 238 | if "Error" in stderr: |
| 239 | print(f"Got error while re-encoding: {stderr}", flush=True) |
| 240 | print("Skipping sample", flush=True) |
| 241 | return None |
| 242 | with io.BytesIO() as video_buffer: |
| 243 | video_buffer.write(stdout_data) |
| 244 | process.wait() |
| 245 | video_buffer.seek(0) |
| 246 | buffer = video_buffer.getvalue() |
| 247 | nr_written_frames = 0 |
| 248 | with io.BytesIO(buffer) as buf, av.open(buf) as container: |
| 249 | for packet in container.demux(): |
| 250 | for frame in packet.decode(): |
| 251 | nr_written_frames = nr_written_frames + 1 |
| 252 | assert nr_written_frames == t, f"Number of frames has changed from {t=} to {nr_written_frames}" |
| 253 | return buffer |
| 254 | |
| 255 | |
| 256 | def wds_filter(sample: dict | None) -> bool: |