From 9816001f9d72ce77ca9f3bebd90b51b32592491c Mon Sep 17 00:00:00 2001 From: nyanpasu64 Date: Sun, 22 Jul 2018 02:06:10 -0700 Subject: [PATCH] [wip] Frame capture and FFmpeg video output --- ovgenpy/outputs.py | 149 +++++++++++++++++++++++++++++++++++++++++ ovgenpy/ovgenpy.py | 34 ++++++++-- ovgenpy/renderer.py | 61 ++++++++++++++--- tests/test_renderer.py | 4 ++ 4 files changed, 233 insertions(+), 15 deletions(-) create mode 100644 ovgenpy/outputs.py diff --git a/ovgenpy/outputs.py b/ovgenpy/outputs.py new file mode 100644 index 0000000..ed8f4d0 --- /dev/null +++ b/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 diff --git a/ovgenpy/ovgenpy.py b/ovgenpy/ovgenpy.py index 69ce010..0c74b80 100644 --- a/ovgenpy/ovgenpy.py +++ b/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 diff --git a/ovgenpy/renderer.py b/ovgenpy/renderer.py index 410b84e..3f88fb8 100644 --- a/ovgenpy/renderer.py +++ b/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.figure import Figure -from matplotlib.lines import Line2D +from matplotlib.backends.backend_agg import FigureCanvasAgg +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.lines: List[Line2D] = None # set by render_frame() first call + self.fig: 'Figure' = None + 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 diff --git a/tests/test_renderer.py b/tests/test_renderer.py index d31cf12..218cd92 100644 --- a/tests/test_renderer.py +++ b/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