[wip] Frame capture and FFmpeg video output

ovgenpy/outputs.py

@@ -0,0 +1,149 @@
+# https://ffmpeg.org/ffplay.html
+import shlex
+import subprocess
+from abc import ABC, abstractmethod
+from io import BytesIO
+from typing import Optional, TYPE_CHECKING, Type, Callable, TypeVar, List
+
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from ovgenpy.ovgenpy import Config
+
+
+IMAGE_FORMAT = 'png'
+
+
+class OutputConfig:
+    cls: 'Type[Output]'
+
+    def __call__(self, ovgen_cfg: 'Config'):
+        return self.cls(ovgen_cfg, cfg=self)
+
+
+class Output(ABC):
+    def __init__(self, ovgen_cfg: 'Config', cfg: OutputConfig):
+        self.ovgen_cfg = ovgen_cfg
+        self.cfg = cfg
+
+    @abstractmethod
+    def output_frame_png(self, frame: bytes) -> None:
+        """ Output an encoded PNG file. TODO PNG compression overhead is bad """
+        pass
+
+
+# Glue logic
+
+def register_output(config_t: Type[OutputConfig]):
+    def inner(output_t: Type[Output]):
+        config_t.cls = output_t
+        return output_t
+
+    return inner
+
+
+# Output subclasses
+
+## FFMPEG templates TODO rename to "...template..."
+FFMPEG = 'ffmpeg'
+FFPLAY = 'ffplay'
+
+
+def ffmpeg_input_video(fps: int) -> List[str]:
+    # Removed: '-c:v {IMAGE_FORMAT}' since it doesn't work
+    return ['-f image2pipe -framerate', str(fps), '-i -']
+
+
+def ffmpeg_input_audio(audio_path: str) -> List[str]:
+    return ['-i', audio_path]
+
+
+FFMPEG_OUTPUT_VIDEO_DEFAULT = '-c:v libx264 -crf 18 -bf 2 -flags +cgop -pix_fmt yuv420p -movflags faststart'
+FFMPEG_OUTPUT_AUDIO_DEFAULT = '-c:a aac -b:a 384k'
+
+
+def parse_templates(templates: List[str]) -> List[str]:
+    return [arg
+            for template in templates
+            for arg in shlex.split(template)]
+
+
+# @dataclass
+# class FFmpegCommand:
+#     audio: Optional[str] = None
+#
+#     def generate_command(self):
+
+
+@dataclass
+class FFmpegOutputConfig(OutputConfig):
+    path: str
+    video_template: str = FFMPEG_OUTPUT_VIDEO_DEFAULT
+    audio_template: str = FFMPEG_OUTPUT_AUDIO_DEFAULT
+
+
+@register_output(FFmpegOutputConfig)
+class FFmpegOutput(Output):
+    # TODO https://github.com/kkroening/ffmpeg-python
+
+    def __init__(self, ovgen_cfg: 'Config', cfg: FFmpegOutputConfig):
+        super().__init__(ovgen_cfg, cfg)
+
+        # Input
+        templates: List[str] = [FFMPEG]
+
+        # TODO factor out "get_ffmpeg_input"... what if wrong abstraction?
+        templates += ffmpeg_input_video(fps=ovgen_cfg.fps)  # video
+        if ovgen_cfg.audio_path:
+            templates += ffmpeg_input_audio(audio_path=ovgen_cfg.audio_path)    # audio
+
+        # Output
+        templates.append(cfg.video_template)  # video
+        if ovgen_cfg.audio_path:
+            templates.append(cfg.audio_template)  # audio
+
+        templates.append(cfg.path)  # output filename
+
+        # Split arguments by words
+        args = parse_templates(templates)
+
+        self._popen = subprocess.Popen(args, stdin=subprocess.PIPE)
+        self._stream = self._popen.stdin
+
+        # Python documentation discourages accessing popen.stdin. It's wrong.
+        # https://stackoverflow.com/a/9886747
+
+    def output_frame_png(self, frame: bytes) -> None:
+        self._stream.write(frame)
+
+    def close(self):
+        self._stream.close()
+        self._popen.wait()
+    # {ffmpeg}
+    #
+    #     # input
+    #     -f image2pipe -framerate {framerate} -c:v {IMAGE_FORMAT} -i {img}
+    #     -i {audio}
+    #
+    #     # output
+    #     -c:a aac -b:a 384k
+    #     -c:v libx264 -crf 18 -bf 2 -flags +cgop -pix_fmt yuv420p -movflags faststart
+    #     {outfile}
+
+
+class FFplayOutputConfig(OutputConfig):
+    pass
+
+@register_output(FFplayOutputConfig)
+class FFplayOutput(Output):
+    pass
+
+
+@dataclass
+class ImageOutputConfig:
+    path_prefix: str
+
+
+@register_output(ImageOutputConfig)
+class ImageOutput(Output):
+    pass

ovgenpy/ovgenpy.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from typing import NamedTuple, Optional, List
 
 import click
+from ovgenpy import outputs
 
 from ovgenpy.renderer import MatplotlibRenderer, RendererConfig
 from ovgenpy.triggers import TriggerConfig, CorrelationTrigger
@@ -23,6 +24,8 @@ class Config(NamedTuple):
 
     trigger: TriggerConfig  # Maybe overriden per Wave
     render: RendererConfig
+    outputs: List[outputs.OutputConfig]
+    create_window: bool
 
     @property
     def time_visible_s(self) -> float:
@@ -37,9 +40,10 @@ _FPS = 60  # f_s
 
 @click.command()
 @click.argument('wave_dir', type=Folder)
-@click.option('--master-wave', type=File, default=None)
+@click.option('--audio_path', type=File, default=None)
 @click.option('--fps', default=_FPS)
-def main(wave_dir: str, audio_path: Optional[str], fps: int):
+@click.option('--output', default='output.mp4')
+def main(wave_dir: str, audio_path: Optional[str], fps: int, output: str):
     cfg = Config(
         wave_dir=wave_dir,
         audio_path=audio_path,
@@ -57,7 +61,11 @@ def main(wave_dir: str, audio_path: Optional[str], fps: int):
         render=RendererConfig(     # todo
             1280, 720,
             ncols=1
-        )
+        ),
+        outputs=[
+            outputs.FFmpegOutputConfig(output)
+        ],
+        create_window=True
     )
 
     ovgen = Ovgen(cfg)
@@ -72,9 +80,11 @@ class Ovgen:
         self.cfg = cfg
         self.waves: List[Wave] = []
         self.nwaves: int = None
+        self.outputs: List[outputs.Output] = []
 
     def write(self):
         self._load_waves()  # self.waves =
+        self._load_outputs()  # self.outputs =
         self._render()
 
     def _load_waves(self):
@@ -97,15 +107,22 @@ class Ovgen:
 
         self.nwaves = len(self.waves)
 
+    def _load_outputs(self):
+        self.outputs = []
+        for output_cfg in self.cfg.outputs:
+            output = output_cfg(self.cfg)
+            self.outputs.append(output)
+
     def _render(self):
         # Calculate number of frames (TODO master file?)
         time_visible_s = self.cfg.time_visible_s
         fps = self.cfg.fps
+        create_window = self.cfg.create_window
 
         nframes = fps * self.waves[0].get_s()
         nframes = int(nframes) + 1
 
-        renderer = MatplotlibRenderer(self.cfg.render, self.nwaves)
+        renderer = MatplotlibRenderer(self.cfg.render, self.nwaves, create_window)
 
         if RENDER_PROFILING:
             begin = time.perf_counter()
@@ -125,9 +142,18 @@ class Ovgen:
 
                 datas.append(wave.get_around(trigger_sample, region_len))
 
+            # Render frame
             print(frame)
             renderer.render_frame(datas)
 
+            # Output frame
+            frame = renderer.get_frame()
+
+            # TODO write to file
+            # how to write ndarray to ffmpeg?
+            # idea: imageio.mimwrite(stdout, ... wait it's blocking = bad
+            # idea: -f rawvideo, pass cfg.render.options... to ffmpeg_input_video()
+
         if RENDER_PROFILING:
             # noinspection PyUnboundLocalVariable
             dtime = time.perf_counter() - begin

ovgenpy/renderer.py

@@ -1,14 +1,19 @@
-from typing import Optional, List, Tuple
+from typing import Optional, List, Tuple, TYPE_CHECKING
 
 import numpy as np
 from dataclasses import dataclass
 from matplotlib import pyplot as plt
-from matplotlib.axes import Axes
+from matplotlib.backends.backend_agg import FigureCanvasAgg
-from matplotlib.figure import Figure
-from matplotlib.lines import Line2D
 
+from ovgenpy.outputs import IMAGE_FORMAT
 from ovgenpy.util import ceildiv
 
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+    from matplotlib.figure import Figure
+    from matplotlib.lines import Line2D
+
+
 
 @dataclass
 class RendererConfig:
@@ -33,6 +38,9 @@ class RendererConfig:
 
 class MatplotlibRenderer:
     """
+    Renderer backend which takes data and produces images.
+    Does not touch Wave or Channel.
+
     If __init__ reads cfg, cfg cannot be hotswapped.
 
     Reasons to hotswap cfg: RendererCfg:
@@ -51,10 +59,10 @@ class MatplotlibRenderer:
 
     DPI = 96
 
-    def __init__(self, cfg: RendererConfig, nplots: int):
+    def __init__(self, cfg: RendererConfig, nplots: int, create_window: bool):
         self.cfg = cfg
         self.nplots = nplots
-        self.fig: Figure = None
+        self.create_window = create_window
 
         # Setup layout
         # "ncols=1" is good for vertical layouts.
@@ -64,8 +72,9 @@ class MatplotlibRenderer:
         self.ncols = 0
 
         # Flat array of nrows*ncols elements, ordered by cfg.rows_first.
-        self.axes: List[Axes] = None        # set by set_layout()
+        self.fig: 'Figure' = None
-        self.lines: List[Line2D] = None     # set by render_frame() first call
+        self.axes: List['Axes'] = None        # set by set_layout()
+        self.lines: List['Line2D'] = None     # set by render_frame() first call
 
         self.set_layout()   # mutates self
 
@@ -85,7 +94,7 @@ class MatplotlibRenderer:
         if self.fig:
             plt.close(self.fig)     # FIXME
 
-        axes2d: np.ndarray[Axes]
+        axes2d: np.ndarray['Axes']
         self.fig, axes2d = plt.subplots(
             self.nrows, self.ncols,
             squeeze=False,
@@ -101,7 +110,7 @@ class MatplotlibRenderer:
         if self.cfg.ncols:
             axes2d = axes2d.T
 
-        self.axes: List[Axes] = axes2d.flatten().tolist()[:self.nplots]
+        self.axes: List['Axes'] = axes2d.flatten().tolist()[:self.nplots]
 
         # Setup figure geometry
         self.fig.set_dpi(self.DPI)
@@ -109,7 +118,8 @@ class MatplotlibRenderer:
             self.cfg.width / self.DPI,
             self.cfg.height / self.DPI
         )
-        plt.show(block=False)
+        if self.create_window:
+            plt.show(block=False)
 
     def _calc_layout(self) -> Tuple[int, int]:
         """
@@ -156,3 +166,32 @@ class MatplotlibRenderer:
 
         self.fig.canvas.draw()
         self.fig.canvas.flush_events()
+
+    assert IMAGE_FORMAT == 'png'
+    RGB_DEPTH = 3
+
+    def get_frame(self):
+        canvas = self.fig.canvas
+
+        # Agg is the default noninteractive backend except on OSX.
+        # https://matplotlib.org/faq/usage_faq.html
+        if not isinstance(canvas, FigureCanvasAgg):
+            raise RuntimeError(
+                f'oh shit, cannot read data from {type(canvas)} != FigureCanvasAgg')
+
+        # buffer_rgba, (w, h) = canvas.print_to_buffer()
+
+        w, h = canvas.get_width_height()
+        buffer_rgb = np.frombuffer(canvas.tostring_rgb(), np.uint8)
+        print(buffer_rgb.shape)
+        np.reshape(buffer_rgb, (w, h, self.RGB_DEPTH))
+
+        return buffer_rgb
+        # # TODO https://matplotlib.org/api/_as_gen/matplotlib.pyplot.imsave.html to
+        # # in-memory stream as png
+        #
+        # # or imsave(arr=...)
+        #
+        # # TODO http://www.icare.univ-lille1.fr/tutorials/convert_a_matplotlib_figure
+        #
+        # raise NotImplementedError

tests/test_renderer.py

@@ -43,3 +43,7 @@ def test_renderer(mock_show):
     # 2 columns, 8 rows
     assert r.ncols == 2
     assert r.nrows == 8
+
+
+# TODO: test get_frame()
+# (integration test) ensure rendering to output works