mirror of
https://git.k-space.ee/arti/vt1200.git
synced 2026-11-01 12:14:56 +02:00
97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
"""Video extraction and the small image primitives the pipeline shares."""
|
|
from __future__ import annotations
|
|
|
|
import subprocess
|
|
from typing import Iterator, Sequence
|
|
|
|
import numpy as np
|
|
|
|
|
|
def temporal_median3(frames: Sequence[np.ndarray]) -> np.ndarray:
|
|
"""Median of three same-shaped frames, as uint8."""
|
|
if len(frames) != 3:
|
|
raise ValueError("temporal_median3 wants exactly three frames")
|
|
return np.median(np.stack(frames), axis=0).astype(np.uint8)
|
|
|
|
|
|
def box_blur3(gray: np.ndarray) -> np.ndarray:
|
|
"""3x3 box blur with edge clamping (no wraparound)."""
|
|
if gray.ndim != 2:
|
|
raise ValueError("box_blur3 wants a 2-D gray image")
|
|
p = np.pad(gray.astype(np.float32), 1, mode="edge")
|
|
total = (
|
|
p[0:-2, 0:-2] + p[0:-2, 1:-1] + p[0:-2, 2:]
|
|
+ p[1:-1, 0:-2] + p[1:-1, 1:-1] + p[1:-1, 2:]
|
|
+ p[2:, 0:-2] + p[2:, 1:-1] + p[2:, 2:]
|
|
)
|
|
return np.rint(total / 9.0).astype(np.uint8)
|
|
|
|
|
|
def binarize(gray: np.ndarray, threshold: int = 127) -> np.ndarray:
|
|
"""True means black ink (the source is dark-on-light)."""
|
|
return gray < threshold
|
|
|
|
|
|
def pack_lsb(mask: np.ndarray) -> bytes:
|
|
"""Pack rows LSB-first: bit 0 of a byte is the leftmost pixel."""
|
|
if mask.ndim != 2 or mask.shape[1] % 8:
|
|
raise ValueError("pack_lsb wants 2-D rows whose width is a multiple of 8")
|
|
return np.packbits(mask.astype(bool), axis=1, bitorder="little").tobytes()
|
|
|
|
|
|
def extract_gray(video: str, start: float, duration: float, fps: int,
|
|
width: int, height: int) -> Iterator[np.ndarray]:
|
|
"""Yield (height, width) uint8 gray frames from ffmpeg."""
|
|
vf = "fps=%s,scale=%d:%d:flags=lanczos,format=gray" % (
|
|
fps, width, height)
|
|
cmd = [
|
|
"ffmpeg", "-v", "error",
|
|
"-ss", str(start),
|
|
"-i", video,
|
|
"-t", str(duration),
|
|
"-vf", vf,
|
|
"-f", "rawvideo", "-",
|
|
]
|
|
size = width * height
|
|
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=size * 4)
|
|
try:
|
|
while True:
|
|
buf = proc.stdout.read(size)
|
|
if len(buf) < size:
|
|
break
|
|
yield np.frombuffer(buf, dtype=np.uint8).reshape(height, width)
|
|
finally:
|
|
proc.stdout.close()
|
|
proc.wait()
|
|
if proc.returncode != 0:
|
|
raise RuntimeError("ffmpeg failed with status %d" % proc.returncode)
|
|
|
|
|
|
def extract_binary(video: str, start: float, duration: float, fps: int,
|
|
width: int, height: int,
|
|
params: dict) -> Iterator[np.ndarray]:
|
|
"""Yield stylised binary frames: temporal median, blur, threshold.
|
|
|
|
Exactly ``duration * fps`` frames are produced. The first output is the
|
|
first frame against itself; each later output is the median of the
|
|
previous, current and next frame.
|
|
"""
|
|
threshold = int(params.get("threshold", 127))
|
|
want_blur = int(params.get("blur", 3))
|
|
it = extract_gray(video, start, duration, fps, width, height)
|
|
try:
|
|
first = next(it)
|
|
except StopIteration:
|
|
return
|
|
|
|
def render(median_frame: np.ndarray) -> np.ndarray:
|
|
gray = box_blur3(median_frame) if want_blur >= 3 else median_frame
|
|
return binarize(gray, threshold)
|
|
|
|
a = first
|
|
b = first
|
|
yield render(temporal_median3([a, b, first]))
|
|
for current in it:
|
|
yield render(temporal_median3([a, b, current]))
|
|
a, b = b, current
|