refactor main.py to support media folders

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]} <file-or-folder>")
+        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()