A simple webserver for serving videos from a directory of dumped frames

frame_server.py

"""
A simple webserver for serving videos dumped as frames as real videos.
Point it towards a file hierarchy that looks like so:
| root
| |- video_1
| |  |- frame_0001.jpg
| |  |- frame_0002.jpg
| |  |- frame_0003.jpg

Video frames are read in 'natural order' by the `natsort` library.

Video are accessed at http://host:port/<video_id>

"""

import argparse
from pathlib import Path
import sys
from PIL import Image
import av
import numpy as np
import io
from flask import abort, send_file, Flask, Response
from natsort import natsorted

parser = argparse.ArgumentParser(description="Serve videos from a directory of directories, each of which has jpg frames in it")
parser.add_argument("root", type=Path)
parser.add_argument("--fps", type=float, default=24)
parser.add_argument("--port", type=int, default=8080)
parser.add_argument("--host", type=str, default="0.0.0.0")
parser.add_argument("--debug", action='store_true')


def reshape_to_even_length_sides(frames):
    if frames.shape[-2] % 2 != 0:
        frames = frames[..., :-1, :]
    if frames.shape[-3] % 2 != 0:
        frames = frames[..., :-1, :, :]
    return frames


def frames_to_video(frames, fps):
    if frames.shape[-1] != 3:
        raise ValueError(f"Expected 3 channel input, but input had shape {frames.shape}")
    # some browsers (e.g. ff) can only decode videos with even length sides
    frames = reshape_to_even_length_sides(frames)
    n_frames, height, width, channels = frames.shape

    with io.BytesIO() as f:
        container = av.open(f, mode='w', format='mp4')
        try:
            stream = container.add_stream('libx264', rate=fps)
            stream.width = width
            stream.height = height
            stream.pix_fmt = 'yuv420p'
            for frame in frames:
                av_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
                for packet in stream.encode(av_frame):
                    container.mux(packet)
            # Flush output
            packet = stream.encode(None)
            container.mux(packet)
        finally:
            container.close()
        return f.getvalue()


def load_frames(directory):
    directory = Path(directory)
    frame_paths = natsorted([p for p in directory.iterdir() if p.suffix.lower() == '.jpg'])
    frames = np.stack([np.array(Image.open(str(p))) for p in frame_paths])
    return frames


def main(argv=sys.argv):
    args = parser.parse_args(argv[1:])
    video_root = args.root
    fps = args.fps
    app = Flask("frame_server")

    @app.route("/<uid>")
    def get_video(uid):
        video_folder_path = video_root / uid
        if video_folder_path.exists():
            frames = load_frames(video_folder_path)
            video_bytes = frames_to_video(frames, fps=fps)
            return send_file(io.BytesIO(video_bytes), mimetype="video/mp4", download_name=f'{uid}.mp4')
        else:
            abort(404)

    app.run(host=args.host, port=args.port, debug=args.debug)

if __name__ == '__main__':
    main()