kopia lustrzana https://github.com/corrscope/corrscope
Merge pull request #25 from nyanpasu64/global-colors
Add support for global colors and per-line colorpull/357/head
nyanpasu64
GitHub
commit
b0e6d827e1
|
|
@ -19,7 +19,6 @@ class ChannelConfig:
|
|||
|
||||
ampl_ratio: float = 1.0 # TODO use amplification = None instead?
|
||||
line_color: Any = None
|
||||
background_color: Any = None
|
||||
|
||||
|
||||
class Channel:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from abc import ABC, abstractmethod
|
|||
from typing import TYPE_CHECKING, Type, List, Union
|
||||
|
||||
from ovgenpy.config import register_config
|
||||
from ovgenpy.utils.keyword_dataclasses import field
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ovgenpy.ovgenpy import Config
|
||||
|
|
@ -100,15 +101,16 @@ class ProcessOutput(Output):
|
|||
# but results in slightly higher CPU consumption.
|
||||
self._stream.write(frame)
|
||||
|
||||
def close(self):
|
||||
def close(self) -> int:
|
||||
self._stream.close()
|
||||
self._popen.wait()
|
||||
return self._popen.wait()
|
||||
|
||||
|
||||
# FFmpegOutput
|
||||
@register_config
|
||||
class FFmpegOutputConfig(IOutputConfig):
|
||||
path: str
|
||||
args: str = ''
|
||||
|
||||
# Do not use `-movflags faststart`, I get corrupted mp4 files (missing MOOV)
|
||||
video_template: str = '-c:v libx264 -crf 18 -preset superfast'
|
||||
|
|
@ -124,6 +126,7 @@ class FFmpegOutput(ProcessOutput):
|
|||
|
||||
ffmpeg = _FFmpegCommand([FFMPEG, '-y'], ovgen_cfg)
|
||||
ffmpeg.add_output(cfg)
|
||||
ffmpeg.templates.append(cfg.args)
|
||||
self.open(ffmpeg.popen([cfg.path], self.bufsize))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ class Ovgen:
|
|||
nframes = int(nframes) + 1
|
||||
|
||||
renderer = MatplotlibRenderer(self.cfg.render, self.cfg.layout, self.nchan)
|
||||
renderer.set_colors(self.cfg.channels)
|
||||
|
||||
if RENDER_PROFILING:
|
||||
begin = time.perf_counter()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
from typing import Optional, List, TYPE_CHECKING, TypeVar, Callable
|
||||
from typing import Optional, List, TYPE_CHECKING, TypeVar, Callable, Any
|
||||
|
||||
import matplotlib
|
||||
import numpy as np
|
||||
|
||||
from ovgenpy.config import register_config
|
||||
from ovgenpy.outputs import RGB_DEPTH
|
||||
from ovgenpy.util import ceildiv
|
||||
from ovgenpy.util import ceildiv, coalesce
|
||||
|
||||
matplotlib.use('agg')
|
||||
from matplotlib import pyplot as plt
|
||||
|
|
@ -15,13 +15,26 @@ if TYPE_CHECKING:
|
|||
from matplotlib.axes import Axes
|
||||
from matplotlib.figure import Figure
|
||||
from matplotlib.lines import Line2D
|
||||
from ovgenpy.channel import ChannelConfig
|
||||
|
||||
|
||||
@register_config
|
||||
def default_color():
|
||||
colors = np.array([int(x, 16) for x in '1f 77 b4'.split()], dtype=float)
|
||||
colors /= np.amax(colors)
|
||||
colors **= 1/3
|
||||
|
||||
return tuple(colors.tolist()) # tolist() converts np.float64 to float
|
||||
|
||||
|
||||
@register_config(always_dump='bg_color init_line_color line_width')
|
||||
class RendererConfig:
|
||||
width: int
|
||||
height: int
|
||||
|
||||
bg_color: Any = 'black'
|
||||
init_line_color: Any = default_color()
|
||||
line_width: Optional[float] = None # TODO
|
||||
|
||||
create_window: bool = False
|
||||
|
||||
|
||||
|
|
@ -54,9 +67,11 @@ class MatplotlibRenderer:
|
|||
self.layout = RendererLayout(lcfg, nplots)
|
||||
|
||||
# Flat array of nrows*ncols elements, ordered by cfg.rows_first.
|
||||
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._fig: 'Figure' = None
|
||||
self._axes: List['Axes'] = None # set by set_layout()
|
||||
self._lines: List['Line2D'] = None # set by render_frame() first call
|
||||
|
||||
self._line_colors: List = [None] * nplots
|
||||
|
||||
self._set_layout() # mutates self
|
||||
|
||||
|
|
@ -71,12 +86,12 @@ class MatplotlibRenderer:
|
|||
|
||||
# Create Axes
|
||||
# https://matplotlib.org/api/_as_gen/matplotlib.pyplot.subplots.html
|
||||
if self.fig:
|
||||
if self._fig:
|
||||
raise Exception("I don't currently expect to call set_layout() twice")
|
||||
plt.close(self.fig)
|
||||
# plt.close(self.fig)
|
||||
|
||||
axes2d: np.ndarray['Axes']
|
||||
self.fig, axes2d = plt.subplots(
|
||||
self._fig, axes2d = plt.subplots(
|
||||
self.layout.nrows, self.layout.ncols,
|
||||
squeeze=False,
|
||||
# Remove gaps between Axes
|
||||
|
|
@ -88,17 +103,29 @@ class MatplotlibRenderer:
|
|||
ax.set_axis_off()
|
||||
|
||||
# Generate arrangement (using nplots, cfg.orientation)
|
||||
self.axes = self.layout.arrange(lambda row, col: axes2d[row, col])
|
||||
self._axes = self.layout.arrange(lambda row, col: axes2d[row, col])
|
||||
|
||||
# Setup figure geometry
|
||||
self.fig.set_dpi(self.DPI)
|
||||
self.fig.set_size_inches(
|
||||
self._fig.set_dpi(self.DPI)
|
||||
self._fig.set_size_inches(
|
||||
self.cfg.width / self.DPI,
|
||||
self.cfg.height / self.DPI
|
||||
)
|
||||
if self.cfg.create_window:
|
||||
plt.show(block=False)
|
||||
|
||||
def set_colors(self, channel_cfgs: List['ChannelConfig']):
|
||||
if len(channel_cfgs) != self.nplots:
|
||||
raise ValueError(
|
||||
f"cannot assign {len(channel_cfgs)} colors to {self.nplots} plots"
|
||||
)
|
||||
|
||||
if self._lines is not None:
|
||||
raise ValueError(
|
||||
f'cannot set line colors after calling render_frame()'
|
||||
)
|
||||
self._line_colors = [cfg.line_color for cfg in channel_cfgs]
|
||||
|
||||
def render_frame(self, datas: List[np.ndarray]) -> None:
|
||||
ndata = len(datas)
|
||||
if self.nplots != ndata:
|
||||
|
|
@ -106,28 +133,36 @@ class MatplotlibRenderer:
|
|||
f'incorrect data to plot: {self.nplots} plots but {ndata} datas')
|
||||
|
||||
# Initialize axes and draw waveform data
|
||||
if self.lines is None:
|
||||
self.lines = []
|
||||
if self._lines is None:
|
||||
self._fig.set_facecolor(self.cfg.bg_color)
|
||||
line_width = self.cfg.line_width
|
||||
|
||||
self._lines = []
|
||||
for idx, data in enumerate(datas):
|
||||
ax = self.axes[idx]
|
||||
# Setup colors
|
||||
line_color = coalesce(self._line_colors[idx], self.cfg.init_line_color)
|
||||
|
||||
# Setup axes
|
||||
ax = self._axes[idx]
|
||||
ax.set_xlim(0, len(data) - 1)
|
||||
ax.set_ylim(-1, 1)
|
||||
|
||||
line = ax.plot(data)[0]
|
||||
self.lines.append(line)
|
||||
# Plot line
|
||||
line = ax.plot(data, color=line_color, linewidth=line_width)[0]
|
||||
self._lines.append(line)
|
||||
|
||||
# Draw waveform data
|
||||
else:
|
||||
for idx, data in enumerate(datas):
|
||||
line = self.lines[idx]
|
||||
line = self._lines[idx]
|
||||
line.set_ydata(data)
|
||||
|
||||
self.fig.canvas.draw()
|
||||
self.fig.canvas.flush_events()
|
||||
self._fig.canvas.draw()
|
||||
self._fig.canvas.flush_events()
|
||||
|
||||
def get_frame(self) -> bytes:
|
||||
""" Returns ndarray of shape w,h,3. """
|
||||
canvas = self.fig.canvas
|
||||
canvas = self._fig.canvas
|
||||
|
||||
# Agg is the default noninteractive backend except on OSX.
|
||||
# https://matplotlib.org/faq/usage_faq.html
|
||||
|
|
|
|||
|
|
@ -8,6 +8,15 @@ def ceildiv(n, d):
|
|||
return -(-n // d)
|
||||
|
||||
|
||||
def coalesce(*args):
|
||||
if len(args) == 0:
|
||||
raise TypeError('coalesce expected 1 argument, got 0')
|
||||
for arg in args:
|
||||
if arg is not None:
|
||||
return arg
|
||||
return args[-1]
|
||||
|
||||
|
||||
T = TypeVar('T')
|
||||
|
||||
# Adapted from https://github.com/numpy/numpy/issues/2269#issuecomment-14436725
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from matplotlib.colors import to_rgb
|
||||
|
||||
from ovgenpy.channel import ChannelConfig
|
||||
from ovgenpy.outputs import RGB_DEPTH
|
||||
from ovgenpy.renderer import RendererConfig, MatplotlibRenderer, LayoutConfig, \
|
||||
RendererLayout
|
||||
|
||||
|
||||
WIDTH = 640
|
||||
HEIGHT = 360
|
||||
|
||||
|
|
@ -90,9 +91,90 @@ def test_renderer():
|
|||
assert r.layout.ncols == 2
|
||||
assert r.layout.nrows == 8
|
||||
|
||||
# TODO: test get_frame()
|
||||
def test_colors():
|
||||
pass # TODO
|
||||
|
||||
ALL_ZEROS = np.array([0,0])
|
||||
|
||||
all_colors = pytest.mark.parametrize('bg_str,fg_str', [
|
||||
('#000000', '#ffffff'),
|
||||
('#ffffff', '#000000'),
|
||||
('#0000aa', '#aaaa00'),
|
||||
('#aaaa00', '#0000aa'),
|
||||
])
|
||||
|
||||
|
||||
@all_colors
|
||||
def test_default_colors(bg_str, fg_str):
|
||||
""" Test the default background/foreground colors. """
|
||||
cfg = RendererConfig(
|
||||
WIDTH,
|
||||
HEIGHT,
|
||||
bg_color=bg_str,
|
||||
init_line_color=fg_str,
|
||||
)
|
||||
lcfg = LayoutConfig()
|
||||
nplots = 1
|
||||
|
||||
r = MatplotlibRenderer(cfg, lcfg, nplots)
|
||||
verify(r, bg_str, fg_str)
|
||||
|
||||
# Ensure default ChannelConfig(line_color=None) does not override line color
|
||||
r = MatplotlibRenderer(cfg, lcfg, nplots)
|
||||
chan = ChannelConfig(wav_path='')
|
||||
r.set_colors([chan] * nplots)
|
||||
verify(r, bg_str, fg_str)
|
||||
|
||||
|
||||
@all_colors
|
||||
def test_line_colors(bg_str, fg_str):
|
||||
""" Test channel-specific line color overrides """
|
||||
cfg = RendererConfig(
|
||||
WIDTH,
|
||||
HEIGHT,
|
||||
bg_color=bg_str,
|
||||
init_line_color='#888888',
|
||||
)
|
||||
lcfg = LayoutConfig()
|
||||
nplots = 1
|
||||
|
||||
r = MatplotlibRenderer(cfg, lcfg, nplots)
|
||||
chan = ChannelConfig(wav_path='', line_color=fg_str)
|
||||
r.set_colors([chan] * nplots)
|
||||
verify(r, bg_str, fg_str)
|
||||
|
||||
|
||||
def verify(r: MatplotlibRenderer, bg_str, fg_str):
|
||||
r.render_frame([ALL_ZEROS])
|
||||
frame_colors: np.ndarray = \
|
||||
np.frombuffer(r.get_frame(), dtype=np.uint8).reshape((-1, RGB_DEPTH))
|
||||
|
||||
bg_u8 = [round(c*255) for c in to_rgb(bg_str)]
|
||||
fg_u8 = [round(c*255) for c in to_rgb(fg_str)]
|
||||
|
||||
# Ensure background is correct
|
||||
assert (frame_colors[0] == bg_u8).all()
|
||||
|
||||
# Ensure foreground is present
|
||||
assert np.prod(frame_colors == fg_u8, axis=-1).any()
|
||||
|
||||
assert (np.amax(frame_colors, axis=0) == np.maximum(bg_u8, fg_u8)).all()
|
||||
assert (np.amin(frame_colors, axis=0) == np.minimum(bg_u8, fg_u8)).all()
|
||||
|
||||
|
||||
# TODO (integration test) ensure rendering to output works
|
||||
|
||||
|
||||
def test_render_output():
|
||||
""" Ensure rendering to output does not raise exceptions. """
|
||||
|
||||
from ovgenpy.ovgenpy import default_config
|
||||
from ovgenpy.outputs import FFmpegOutput, FFmpegOutputConfig
|
||||
|
||||
cfg = default_config(render=RendererConfig(WIDTH, HEIGHT))
|
||||
renderer = MatplotlibRenderer(cfg.render, cfg.layout, nplots=1)
|
||||
output_cfg = FFmpegOutputConfig('-', '-f nut')
|
||||
out = FFmpegOutput(cfg, output_cfg)
|
||||
|
||||
renderer.render_frame([ALL_ZEROS])
|
||||
out.write_frame(renderer.get_frame())
|
||||
|
||||
assert out.close() == 0
|
||||
|
|
|
|||
Ładowanie…
Reference in New Issue