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:])

examples/streamer/samples/generate_opus.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+
+from kaitaistruct import KaitaiStruct, ValidationNotEqualError
+import os
+import getopt
+import sys
+import glob
+from functools import reduce
+
+
+class Ogg(KaitaiStruct):
+    """Ogg is a popular media container format, which provides basic
+    streaming / buffering mechanisms and is content-agnostic. Most
+    popular codecs that are used within Ogg streams are Vorbis (thus
+    making Ogg/Vorbis streams) and Theora (Ogg/Theora).
+
+    Ogg stream is a sequence Ogg pages. They can be read sequentially,
+    or one can jump into arbitrary stream location and scan for "OggS"
+    sync code to find the beginning of a new Ogg page and continue
+    decoding the stream contents from that one.
+    """
+
+    def __init__(self, _io, _parent=None, _root=None):
+        KaitaiStruct.__init__(self, _io)
+        self._parent = _parent
+        self._root = _root if _root else self
+        self._read()
+
+    def _read(self):
+        self.pages = []
+        i = 0
+        while not self._io.is_eof():
+            self.pages.append(Ogg.Page(self._io, self, self._root))
+            i += 1
+
+    class Page(KaitaiStruct):
+        """Ogg page is a basic unit of data in an Ogg bitstream, usually
+        it's around 4-8 KB, with a maximum size of 65307 bytes.
+        """
+
+        def __init__(self, _io, _parent=None, _root=None):
+            KaitaiStruct.__init__(self, _io)
+            self._parent = _parent
+            self._root = _root if _root else self
+            self._read()
+
+        def _read(self):
+            self.sync_code = self._io.read_bytes(4)
+            if not self.sync_code == b"\x4F\x67\x67\x53":
+                raise ValidationNotEqualError(b"\x4F\x67\x67\x53", self.sync_code, self._io,
+                                              u"/types/page/seq/0")
+            self.version = self._io.read_bytes(1)
+            if not self.version == b"\x00":
+                raise ValidationNotEqualError(b"\x00", self.version, self._io, u"/types/page/seq/1")
+            self.reserved1 = self._io.read_bits_int_be(5)
+            self.is_end_of_stream = self._io.read_bits_int_be(1) != 0
+            self.is_beginning_of_stream = self._io.read_bits_int_be(1) != 0
+            self.is_continuation = self._io.read_bits_int_be(1) != 0
+            self._io.align_to_byte()
+            self.granule_pos = self._io.read_u8le()
+            self.bitstream_serial = self._io.read_u4le()
+            self.page_seq_num = self._io.read_u4le()
+            self.crc32 = self._io.read_u4le()
+            self.num_segments = self._io.read_u1()
+            self.len_segments = [None] * self.num_segments
+            for i in range(self.num_segments):
+                self.len_segments[i] = self._io.read_u1()
+
+            self.segments = [None] * self.num_segments
+            for i in range(self.num_segments):
+                self.segments[i] = self._io.read_bytes(self.len_segments[i])
+
+
+def generate(input_file: str, output_dir: str, max_samples: int):
+    if output_dir[-1] != "/":
+        output_dir += "/"
+    if os.path.isdir(output_dir):
+        files_to_delete = glob.glob(output_dir + "*.opus")
+        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)
+    audio_stream_file = "_audio_stream.ogg"
+    if os.path.isfile(audio_stream_file):
+        os.remove(audio_stream_file)
+    os.system('ffmpeg -i {} -vn -ar 48000 -ac 2 -vbr off -acodec libopus -ab 64k {}'.format(input_file, audio_stream_file))
+
+    data = Ogg.from_file(audio_stream_file)
+    index = 0
+    valid_pages = data.pages[2:]
+    segments = list(reduce(lambda x, y: x + y.segments, valid_pages, []))[:max_samples]
+    for segment in segments:
+        name = "{}sample-{}.opus".format(output_dir, index)
+        index += 1
+        with open(name, 'wb') as file:
+            assert len(list(segment)) == 160
+            file.write(segment)
+    os.remove(audio_stream_file)
+
+
+def main(argv):
+    input_file = None
+    default_output_dir = "opus/"
+    output_dir = default_output_dir
+    max_samples = None
+    try:
+        opts, args = getopt.getopt(argv, "hi:o:m:", ["help", "ifile=", "odir=", "max="])
+    except getopt.GetoptError:
+        print('generate_opus.py -i <input_files> [-o <output_files>] [-m <max_samples>] [-h]')
+        sys.exit(2)
+    for opt, arg in opts:
+        if opt in ("-h", "--help"):
+            print("Usage: generate_opus.py -i <input_files> [-o <output_files>] [-m <max_samples>] [-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-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)
+    if input_file is None:
+        print("Missing argument -i")
+        sys.exit(2)
+    generate(input_file, output_dir, max_samples)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])