images => video

app.py

@@ -13,8 +13,6 @@
 # - x: 128 to 512
 # - y: 128 to 512
 # - physical fields: 2 to 8 (density, pressure, energy, velocity…)
-from functools import lru_cache
-
 import gradio as gr
 import h5py
 import numpy as np
@@ -22,8 +20,16 @@ from fsspec import url_to_fs
 from matplotlib import cm
 from PIL import Image
 import av
-import io
+from tempfile import gettempdir
+import os
+
+# Get the path of the system's temporary directory
+temp_directory = gettempdir()
+print(f"System's temporary directory is: {temp_directory}")
+videos_temp_directory = os.path.join(temp_directory, "videos")
+print(f"Videos are saved (and never deleted) in: {videos_temp_directory}")
 
+# TODO: add colormap input
 
 repo_id = "lhoestq/turbulent_radiative_layer_tcool_demo"
 set_path = f"hf://datasets/{repo_id}/**/*.hdf5"
@@ -31,103 +37,127 @@ fs, _ = url_to_fs(set_path)
 paths = fs.glob(set_path)
 files = {path: h5py.File(fs.open(path, "rb", cache_type="none"), "r") for path in paths}
 
+
 def get_scalar_fields(path: str) -> list[str]:
     # TODO: support t1_fields (vector) and t2_fields (tensor)
     return list(files[path]["t0_fields"].keys())
 
+
 def get_trajectories(path: str, field: str) -> list[int]:
     # The first dimension is the trajectory (8 to 256)
     return list(range(len(files[path]["t0_fields"][field])))
 
-@lru_cache(maxsize=4)
-def get_images(path: str, scalar_field: str, trajectory: int) -> list[Image.Image]:
-    # The data is of shape (n_trajectories, n_time_steps, x, y)
-    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)
-    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'
 
+
+def create_video(
+    path: str, scalar_field: str, trajectory: int, video_filename: str
+) -> None:
     out = files[path]["t0_fields"][scalar_field][trajectory]
-    out = np.log(out) # not sure why
+    # 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]
+    out = np.uint8(cm.viridis(out) * 255)
+
+    output = av.open(video_filename, "w")
+    stream = output.add_stream("libvpx-vp9", str(fps))
+    height, width = out[0].shape[1], out[0].shape[0]
     stream.width = width
     stream.height = height
-    stream.pix_fmt = 'yuv444p' # or yuva420p
-    # stream.options = {'crf': '17'}
+    stream.pix_fmt = "yuv444p"
 
     for img in out:
-      image = Image.fromarray(img)
-      frame = av.VideoFrame.from_image(image)
-      packet = stream.encode(frame)
-      output.mux(packet)
+        image = Image.fromarray(img)
+        # I think it's the way to get the expected orientation
+        image = image.transpose(method=Image.Transpose.TRANSPOSE)
+        image = image.transpose(method=Image.Transpose.FLIP_TOP_BOTTOM)
+        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])
+# no limit on the size of the videos on the disk
+def get_video(path: str, scalar_field: str, trajectory: int) -> str:
+    video_filename = os.path.join(
+        videos_temp_directory, *path.split("/"), scalar_field, f"{trajectory}.webm"
+    )
+    os.makedirs(os.path.dirname(video_filename), exist_ok=True)
+    if not os.path.isfile(video_filename):
+        create_video(path, scalar_field, trajectory, video_filename)
+    return video_filename
+
 
 with gr.Blocks() as demo:
-    gr.Markdown(f"# 💠 HDF5 Viewer for the [{repo_id}](https://huggingface.co/datasets/{repo_id}) Dataset 🌊")
+    default_scalar_fields = get_scalar_fields(paths[0])
+    default_trajectories = get_trajectories(paths[0], default_scalar_fields[0])
+    default_video = get_video(
+        paths[0], default_scalar_fields[0], default_trajectories[0]
+    )
+
+    gr.Markdown(
+        f"# 💠 HDF5 Viewer for the [{repo_id}](https://huggingface.co/datasets/{repo_id}) Dataset 🌊"
+    )
     gr.Markdown(f"Showing files at `{set_path}`")
     with gr.Row():
-        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=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, video])
+        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",
+        )
+    video = gr.Video(default_video, height=400)
+
+    @files_dropdown.select(
+        inputs=[files_dropdown],
+        outputs=[scalar_fields_dropdown, trajectory_dropdown, 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),
-            video: gr.Video(vid)
+            scalar_fields_dropdown: gr.Dropdown(
+                choices=scalar_fields, value=scalar_fields[0]
+            ),
+            trajectory_dropdown: gr.Dropdown(
+                choices=trajectories, value=trajectories[0]
+            ),
+            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, video])
+
+    @scalar_fields_dropdown.select(
+        inputs=[files_dropdown, scalar_fields_dropdown],
+        outputs=[trajectory_dropdown, 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),
-            video: gr.Video(vid)
+            trajectory_dropdown: gr.Dropdown(
+                choices=trajectories, value=trajectories[0]
+            ),
+            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, video])
+    @trajectory_dropdown.select(
+        inputs=[files_dropdown, scalar_fields_dropdown, trajectory_dropdown],
+        outputs=[video],
+    )
     def _update_trajectory(path: str, scalar_field: str, trajectory: int):
-        images = get_images(path, scalar_field, trajectory)
         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)}
+        yield {video: gr.Video(vid)}
+
 
-demo.launch()
+demo.launch()