Add python scripts to generate h264 and opus samples

examples/streamer/samples/generate_h264.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+
+import os
+import getopt
+import sys
+import glob
+from functools import reduce
+from typing import Optional, List
+
+
+class H264ByteStream:
+    @staticmethod
+    def nalu_type(nalu: bytes) -> int:
+        return nalu[0] & 0x1F
+
+    @staticmethod
+    def merge_sample(sample: List[bytes]) -> bytes:
+        result = bytes()
+        for nalu in sample:
+            result += len(nalu).to_bytes(4, byteorder='big') + nalu
+        return result
+
+    @staticmethod
+    def reduce_nalus_to_samples(samples: List[List[bytes]], current: bytes) -> List[List[bytes]]:
+        last_nalus = samples[-1]
+        samples[-1] = last_nalus + [current]
+        if H264ByteStream.nalu_type(current) in [1, 5]:
+            samples.append([])
+        return samples
+
+    def __init__(self, file_name: str):
+        with open(file_name, "rb") as file:
+            byte_stream = file.read()
+            long_split = byte_stream.split(b"\x00\x00\x00\x01")
+            splits = reduce(lambda acc, x: acc + x.split(b"\x00\x00\x01"), long_split, [])
+            nalus = filter(lambda x: len(x) > 0, splits)
+            self.samples = list(
+                filter(lambda x: len(x) > 0, reduce(H264ByteStream.reduce_nalus_to_samples, nalus, [[]])))
+
+
+def generate(input_file: str, output_dir: str, max_samples: Optional[int], fps: Optional[int]):
+    if output_dir[-1] != "/":
+        output_dir += "/"
+    if os.path.isdir(output_dir):
+        files_to_delete = glob.glob(output_dir + "*.h264")
+        if len(files_to_delete) > 0:
+            print("Remove following files?")
+            for file in files_to_delete:
+                print(file)
+            response = input("Remove files? [y/n] ").lower()
+            if response != "y" and response != "yes":
+                print("Cancelling...")
+                return
+            print("Removing files")
+            for file in files_to_delete:
+                os.remove(file)
+    else:
+        os.makedirs(output_dir, exist_ok=True)
+    video_stream_file = "_video_stream.h264"
+    if os.path.isfile(video_stream_file):
+        os.remove(video_stream_file)
+
+    fps_line = "" if fps is None else "-filter:v fps=fps={} ".format(fps)
+    command = 'ffmpeg -i {} -an -vcodec libx264 -preset slow -profile baseline {}{}'.format(input_file, fps_line,
+                                                                                            video_stream_file)
+    os.system(command)
+
+    data = H264ByteStream(video_stream_file)
+    index = 0
+    for sample in data.samples[:max_samples]:
+        name = "{}sample-{}.h264".format(output_dir, index)
+        index += 1
+        with open(name, 'wb') as file:
+            merged_sample = H264ByteStream.merge_sample(sample)
+            file.write(merged_sample)
+    os.remove(video_stream_file)
+
+
+def main(argv):
+    input_file = None
+    default_output_dir = "h264/"
+    output_dir = default_output_dir
+    max_samples = None
+    fps = None
+    try:
+        opts, args = getopt.getopt(argv, "hi:o:m:f:", ["help", "ifile=", "odir=", "max=", "fps"])
+    except getopt.GetoptError:
+        print('generate_h264.py -i <input_files> [-o <output_files>] [-m <max_samples>] [-f <fps>] [-h]')
+        sys.exit(2)
+    for opt, arg in opts:
+        if opt in ("-h", "--help"):
+            print("Usage: generate_h264.py -i <input_files> [-o <output_files>] [-m <max_samples>] [-f <fps>] [-h]")
+            print("Arguments:")
+            print("\t-i,--ifile: Input file")
+            print("\t-o,--odir: Output directory (default: " + default_output_dir + ")")
+            print("\t-m,--max: Maximum generated samples")
+            print("\t-f,--fps: Output fps")
+            print("\t-h,--help: Print this help and exit")
+            sys.exit()
+        elif opt in ("-i", "--ifile"):
+            input_file = arg
+        elif opt in ("-o", "--ofile"):
+            output_dir = arg
+        elif opt in ("-m", "--max"):
+            max_samples = int(arg)
+        elif opt in ("-f", "--fps"):
+            fps = int(arg)
+    if input_file is None:
+        print("Missing argument -i")
+        sys.exit(2)
+    generate(input_file, output_dir, max_samples, fps)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])