test create a video

app.py

@@ -21,6 +21,9 @@ import numpy as np
 from fsspec import url_to_fs
 from matplotlib import cm
 from PIL import Image
+import av
+import io
+
 
 repo_id = "lhoestq/turbulent_radiative_layer_tcool_demo"
 set_path = f"hf://datasets/{repo_id}/**/*.hdf5"
@@ -45,9 +48,43 @@ def get_images(path: str, scalar_field: str, trajectory: int) -> list[Image.Imag
     out = np.uint8(cm.RdBu_r(out) * 255)
     return [Image.fromarray(img) for img in out]
 
+fps = 25
+# @lru_cache(maxsize=4)
+def get_video(path: str, scalar_field: str, trajectory: int) -> str:
+    video_filename = 'output_vid.webm'
+
+    out = files[path]["t0_fields"][scalar_field][trajectory]
+    out = np.log(out) # not sure why
+    out = (out - out.min()) / (out.max() - out.min())
+    out = np.uint8(cm.RdBu_r(out) * 255)
+    
+    output = av.open(video_filename, 'w')
+    stream = output.add_stream('libvpx-vp9', str(fps))
+    width, height = out[0].shape[1], out[0].shape[0]
+    stream.width = width
+    stream.height = height
+    stream.pix_fmt = 'yuv444p' # or yuva420p
+    # stream.options = {'crf': '17'}
+
+    for img in out:
+      image = Image.fromarray(img)
+      frame = av.VideoFrame.from_image(image)
+      packet = stream.encode(frame)
+      output.mux(packet)
+
+    # Flush the encoder and close the "in memory" file:
+    packet = stream.encode(None)
+    output.mux(packet)
+    output.close()
+    return video_filename
+    
+    # subprocess.run(["ffmpeg", "-y", "-framerate", "25", "-i", os.path.join(output_dir, "density_%d.png"), "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p", os.path.join(output_dir, "density.webm")])
+
+
 default_scalar_fields = get_scalar_fields(paths[0])
 default_trajectories = get_trajectories(paths[0], default_scalar_fields[0])
 default_images = get_images(paths[0], default_scalar_fields[0], default_trajectories[0])
+default_video = get_video(paths[0], default_scalar_fields[0], default_trajectories[0])
 
 with gr.Blocks() as demo:
     gr.Markdown(f"# 💠 HDF5 Viewer for the [{repo_id}](https://huggingface.co/datasets/{repo_id}) Dataset 🌊")
@@ -56,35 +93,41 @@ with gr.Blocks() as demo:
         files_dropdown = gr.Dropdown(choices=paths, value=paths[0], label="File", scale=4)
         scalar_fields_dropdown = gr.Dropdown(choices=default_scalar_fields, value=default_scalar_fields[0], label="Physical field")
         trajectory_dropdown = gr.Dropdown(choices=default_trajectories, value=default_trajectories[0], label="Trajectory")
-    gallery = gr.Gallery(default_images, preview=True, selected_index=len(default_images) // 2)
+    gallery = gr.Gallery(default_images, preview=False, selected_index=len(default_images) // 2)
     gr.Markdown("_Tip: click on the image to go forward or backwards_")
+    video = gr.Video(default_video)
 
-    @files_dropdown.select(inputs=[files_dropdown], outputs=[scalar_fields_dropdown, trajectory_dropdown, gallery])
+    @files_dropdown.select(inputs=[files_dropdown], outputs=[scalar_fields_dropdown, trajectory_dropdown, gallery, video])
     def _update_file(path: str):
         scalar_fields = get_scalar_fields(path)
         trajectories = get_trajectories(path, scalar_fields[0])
         images = get_images(path, scalar_fields[0], trajectories[0])
+        vid = get_video(path, scalar_fields[0], trajectories[0])
         yield {
             scalar_fields_dropdown: gr.Dropdown(choices=scalar_fields, value=scalar_fields[0]),
             trajectory_dropdown: gr.Dropdown(choices=trajectories, value=trajectories[0]),
-            gallery: gr.Gallery(images)
+            gallery: gr.Gallery(images),
+            video: gr.Video(vid)
         }
         yield {gallery: gr.Gallery(selected_index=len(default_images) // 2)}
     
-    @scalar_fields_dropdown.select(inputs=[files_dropdown, scalar_fields_dropdown], outputs=[trajectory_dropdown, gallery])
+    @scalar_fields_dropdown.select(inputs=[files_dropdown, scalar_fields_dropdown], outputs=[trajectory_dropdown, gallery, video])
     def _update_scalar_field(path: str, scalar_field: str):
         trajectories = get_trajectories(path, scalar_field)
         images = get_images(path, scalar_field, trajectories[0])
+        vid = get_video(path, scalar_field, trajectories[0])
         yield {
             trajectory_dropdown: gr.Dropdown(choices=trajectories, value=trajectories[0]),
-            gallery: gr.Gallery(images)
+            gallery: gr.Gallery(images),
+            video: gr.Video(vid)
         }
         yield {gallery: gr.Gallery(selected_index=len(default_images) // 2)}
 
-    @trajectory_dropdown.select(inputs=[files_dropdown, scalar_fields_dropdown, trajectory_dropdown], outputs=[gallery])
+    @trajectory_dropdown.select(inputs=[files_dropdown, scalar_fields_dropdown, trajectory_dropdown], outputs=[gallery, video])
     def _update_trajectory(path: str, scalar_field: str, trajectory: int):
         images = get_images(path, scalar_field, trajectory)
-        yield {gallery: gr.Gallery(images)}
+        vid = get_video(path, scalar_field, trajectory)
+        yield {gallery: gr.Gallery(images), video: gr.Video(vid)}
         yield {gallery: gr.Gallery(selected_index=len(default_images) // 2)}
 
 demo.launch()

requirements.txt

@@ -2,3 +2,4 @@ h5py
 huggingface_hub
 Pillow
 numpy
+av