willprice · July 6, 2021 10:11 · Jul 6, 2021
diff --git a/frame_server.py b/frame_server.py
@@ -0,0 +1,94 @@
+"""
+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()
No results found