forked from hkc/cc-stuff
1
0
Fork 0
cc-stuff/mess/package_video.py

151 lines
5.0 KiB
Python

# x-run: python3 % ~/downloads/moZtoMP7HAA.mp4 /tmp/video.cani
from typing import Literal
from dataclasses import dataclass
from struct import pack
from sys import argv
from subprocess import Popen, PIPE, run
from tqdm import tqdm
from tempfile import TemporaryDirectory
from glob import glob
from PIL import Image
from functools import lru_cache
PIX_BITS = [[1, 2], [4, 8], [16, 0]]
@dataclass
class VideoMetadata:
framerate: Literal[20, 10, 5] = 10
audio_channels: Literal[1, 2] = 2
sample_rate: Literal[12000, 24000, 48000] = 48000
screen_width: int = 164
screen_height: int = 81
@property
def audio_samples_per_frame(self) -> int:
return self.sample_rate // self.framerate
def serialize(self) -> bytes:
return bytes([ self.framerate, self.audio_channels ]) \
+ pack("<HHH", self.sample_rate, self.screen_width, self.screen_height)
@dataclass
class VideoFrame:
audio: list[bytes]
video: list[bytes]
palette: list[int]
def serialize(self) -> bytes:
return bytes.join(b"", self.audio) \
+ bytes.join(b"", self.video) \
+ bytes.join(b"", [ bytes.fromhex("%06x" % color) for color in self.palette ])
@lru_cache
def _brightness(palette: tuple, i: int) -> float:
r, g, b = palette[i * 3 : (i + 1) * 3]
return (r + g + b) / 768
@lru_cache
def _distance(palette: tuple, a: int, b: int) -> float:
r1, g1, b1 = palette[a * 3 : (a + 1) * 3]
r2, g2, b2 = palette[b * 3 : (b + 1) * 3]
rd, gd, bd = r1 - r2, g1 - g2, b1 - b2
return (rd * rd + gd * gd + bd * bd) / 1966608
@lru_cache
def _get_colors(imgdata, palette: tuple, x: int, y: int) -> tuple[int, int]:
brightest_i, brightest_l = 0, 0
darkest_i, darkest_l = 0, 768
for oy, line in enumerate(PIX_BITS):
for ox in range(len(line)):
pix = imgdata[x + ox, y + oy]
assert pix < 16, f"{pix} is too big at {x+ox}:{y+oy}"
brightness = _brightness(palette, pix)
if brightness > brightest_l:
brightest_l, brightest_i = brightness, pix
if brightness < darkest_l:
darkest_l, darkest_i = brightness, pix
return darkest_i, brightest_i
@lru_cache()
def _is_darker(palette: tuple, bg: int, fg: int, c: int) -> bool:
return _distance(palette, bg, c) < _distance(palette, fg, c)
def _get_block(imgdata, palette: tuple, x: int, y: int) -> tuple[int, int, int]:
dark_i, bri_i = _get_colors(imgdata, palette, x, y)
assert dark_i < 16, f"{dark_i} is too big"
assert bri_i < 16, f"{bri_i} is too big"
out: int = 0
for oy, line in enumerate(PIX_BITS):
for ox, bit in enumerate(line):
if not _is_darker(
palette, dark_i, bri_i, imgdata[x + ox, y + oy]
):
out |= bit
# bottom right pixel fix?
if not _is_darker(palette, dark_i, bri_i, imgdata[x + 1, y + 2]):
out ^= 31
dark_i, bri_i = bri_i, dark_i
return out, dark_i, bri_i
metadata = VideoMetadata(
framerate=20,
audio_channels=2,
sample_rate=24000,
screen_width=164,
screen_height=81
)
input_video = argv[1]
with TemporaryDirectory() as tmpdir:
run([
"ffmpeg",
"-i", input_video,
"-f", "s8",
"-ac", str(metadata.audio_channels),
"-ar", str(metadata.sample_rate),
f"{tmpdir}/audio.s8"
])
run([
"ffmpeg",
"-i", input_video,
"-an",
"-r", str(metadata.framerate),
"-vf", f"scale={metadata.screen_width * 2}:{metadata.screen_height * 3}",
f"{tmpdir}/video%06d.jpg"
])
with open(argv[2], "w") as fp_out, open(f"{tmpdir}/audio.s8", "rb") as fp_audio:
print(metadata.serialize().hex(), file=fp_out)
for i, frame_path in tqdm(enumerate(glob(f"{tmpdir}/video*.jpg"))):
with Image.open(frame_path) as img_in:
img_in = img_in.convert("P", palette=Image.Palette.ADAPTIVE, colors=16)
img_data = img_in.load()
img_palette = tuple(img_in.getpalette()) # type: ignore
audio_samples = fp_audio.read(metadata.audio_samples_per_frame * metadata.audio_channels)
frame = VideoFrame(
[
audio_samples[i::metadata.audio_channels]
for i in range(metadata.audio_channels)
],
[],
[
(r << 16) | (g << 8) | b
for r, g, b
in zip(img_palette[0::3], img_palette[1::3], img_palette[2::3]) # type: ignore
])
for y in range(0, img_in.height - 2, 3):
line = bytearray()
for x in range(0, img_in.width - 1, 2):
ch, bg, fg = _get_block(img_data, img_palette, x, y)
line.extend([(ch + 0x80) & 0xFF, fg << 4 | bg])
frame.video.append(line)
print(frame.serialize().hex(), file=fp_out)