mirror
/
corrscope
kopia lustrzana https://github.com/corrscope/corrscope
1
0
Forkuj 0
Kod Zgłoszenia Packages Projekty Wydania Wiki Aktywność

[wip] Frame capture and FFmpeg video output

pull/357/head
nyanpasu64 2018-07-22 02:06:10 -07:00
rodzic 46df1b8adc
commit 9816001f9d
4 zmienionych plików z 233 dodań i 15 usunięć

149
ovgenpy/outputs.py 100644
Wyświetl plik

@ -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

34
ovgenpy/ovgenpy.py
Wyświetl plik

@ -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

61
ovgenpy/renderer.py
Wyświetl plik

@ -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

4
tests/test_renderer.py
Wyświetl plik

@ -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