diff --git a/main.py b/main.py index 5182284..2477bd9 100644 --- a/main.py +++ b/main.py @@ -1,26 +1,148 @@ -import ffmpeg import os -in_filename = "/home/albmj/media/Cowboy_Bebop.mkv" -width=640 -height=360 -process1 = ( - ffmpeg - .input(in_filename,re=None,stream_loop="-1") - .filter('scale', width, height) - .filter('subtitles', in_filename, force_style='Fontsize=120', si='2') - .output('pipe:', format='rawvideo', pix_fmt='rgb24') - .run_async(pipe_stdout=True,quiet=False) -) -fd = os.open("/tmp/video",os.O_RDWR|os.O_CREAT) -os.lseek(fd,0,os.SEEK_SET) -#fd = os.memfd_create('test_file') -print(f'{width}x{height}x3={width*height*3}') -print(f'/tmp/video') -while True: - in_bytes = process1.stdout.read(width * height * 3) - if len(in_bytes) != width*height*3: - print("wrong size") - print(len(in_bytes)) - exit(1) - os.write(fd, in_bytes) +import signal +import sys +import threading +from types import FrameType +from typing import Any, Literal, cast + +import ffmpeg # type: ignore[import-untyped] + +_ffmpeg: Any = ffmpeg + +WIDTH = 640 +HEIGHT = 360 +FRAME_SIZE = WIDTH * HEIGHT * 3 +OUTPUT = "/tmp/video" +IMAGE_HOLD_SECONDS = 10 + +# Set when the user asks us to shut down (Ctrl-C / SIGTERM). +_shutdown = threading.Event() + +IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp", ".tiff", ".tif"} +VIDEO_EXTS = {".mkv", ".mp4", ".mov", ".avi", ".webm", ".m4v", ".flv", ".wmv", ".mpg", ".mpeg"} + +Kind = Literal["image", "video"] + + +def classify(path: str) -> Kind | None: + ext = os.path.splitext(path)[1].lower() + if ext in IMAGE_EXTS: + return "image" + if ext in VIDEO_EXTS: + return "video" + return None + + +def collect_media(input_path: str) -> list[tuple[str, Kind]]: + """Return a sorted list of (path, kind) for the given file or folder.""" + if os.path.isfile(input_path): + kind = classify(input_path) + if kind is None: + print(f"unsupported file type: {input_path}") + sys.exit(1) + return [(input_path, kind)] + + if os.path.isdir(input_path): + media: list[tuple[str, Kind]] = [] + for name in sorted(os.listdir(input_path)): + full = os.path.join(input_path, name) + if not os.path.isfile(full): + continue + kind = classify(full) + if kind is not None: + media.append((full, kind)) + if not media: + print(f"no images or videos found in: {input_path}") + sys.exit(1) + return media + + print(f"no such file or folder: {input_path}") + sys.exit(1) + + +def write_frame(fd: int, frame: bytes) -> None: os.lseek(fd, 0, os.SEEK_SET) + os.write(fd, frame) + os.lseek(fd, 0, os.SEEK_SET) + + +def show_image(fd: int, path: str) -> None: + """Decode a single image, scale to the display, and hold it for a while.""" + print(f"image: {path}") + stream: Any = ( + _ffmpeg + .input(path) + .filter("scale", WIDTH, HEIGHT) + .output("pipe:", format="rawvideo", pix_fmt="rgb24", vframes=1) + ) + result = cast("tuple[bytes, bytes]", stream.run(capture_stdout=True, quiet=True)) + out = result[0] + if len(out) < FRAME_SIZE: + print(f"skipping {path}: decoded {len(out)} bytes, expected {FRAME_SIZE}") + return + write_frame(fd, out[:FRAME_SIZE]) + # Interruptible hold: returns immediately once shutdown is requested. + _shutdown.wait(IMAGE_HOLD_SECONDS) + + +def play_video(fd: int, path: str) -> None: + """Stream a video through to the end at real-time pace.""" + print(f"video: {path}") + stream: Any = ( + _ffmpeg + .input(path, re=None) + .filter("scale", WIDTH, HEIGHT) + .output("pipe:", format="rawvideo", pix_fmt="rgb24") + ) + process: Any = stream.run_async(pipe_stdout=True, quiet=False) + stdout: Any = process.stdout + try: + while not _shutdown.is_set(): + frame = cast("bytes", stdout.read(FRAME_SIZE)) + if len(frame) != FRAME_SIZE: + break + write_frame(fd, frame) + finally: + # Stop ffmpeg so it doesn't keep decoding after we stop reading. + process.terminate() + stdout.close() + process.wait() + + +def _request_shutdown(signum: int, frame: FrameType | None) -> None: + print("\nshutting down...") + _shutdown.set() + + +def main() -> None: + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ") + sys.exit(1) + + media = collect_media(sys.argv[1]) + + signal.signal(signal.SIGINT, _request_shutdown) + signal.signal(signal.SIGTERM, _request_shutdown) + + fd = os.open(OUTPUT, os.O_RDWR | os.O_CREAT) + # Ensure the file is at least one frame long so the C side can mmap it. + os.ftruncate(fd, FRAME_SIZE) + + print(f"{WIDTH}x{HEIGHT}x3={FRAME_SIZE}") + print(OUTPUT) + + try: + while not _shutdown.is_set(): + for path, kind in media: + if _shutdown.is_set(): + break + if kind == "image": + show_image(fd, path) + else: + play_video(fd, path) + finally: + os.close(fd) + + +if __name__ == "__main__": + main()