Update app.py

app.py

@@ -84,45 +84,35 @@ def decode_frame(frame):
         raise RuntimeError("Failed to decode JPEG frame")
     return img
 
-
 def convert_zarr_from_hf(uuid: str, tmp_dir: str) -> str:
-    """
-    Open images.front_1 directly from HF via fsspec, convert to mp4.
-    Returns path to the output mp4.
-    """
-    fs = HfFileSystem()
-    store_path = f"hf://datasets/{EGOVERSE_REPO}/{uuid}/images.front_1"
+    from huggingface_hub import snapshot_download
 
-    store = zarr.storage.FsspecStore.from_url(store_path)
-    root  = zarr.open(store, mode="r", zarr_format=3)
+    yield_status = f"Downloading zarr for {uuid} ..."  # caller will show this
 
-    if "c" not in root and not hasattr(root, "shape"):
-        # root IS the array when opened at images.front_1
-        frames = root
-    else:
-        frames = root  # images.front_1 is the array itself
+    # Download only the images.front_1 array for this episode
+    local_dir = snapshot_download(
+        repo_id=EGOVERSE_REPO,
+        repo_type="dataset",
+        allow_patterns=[f"{uuid}/images.front_1/**"],
+        local_dir=os.path.join(tmp_dir, uuid),
+    )
+
+    array_path = os.path.join(local_dir, uuid, "images.front_1")
+    root = zarr.open(array_path, mode="r", zarr_format=3)
+    frames = root
 
     n_frames = frames.shape[0]
-    first    = decode_frame(frames[0])
+    first    = decode_frame(frames[0:1][0])
     h, w     = first.shape[:2]
 
+    all_frames = frames[0:n_frames]  # single read, now local so fast
+
     temp_avi = os.path.join(tmp_dir, f"{uuid}.avi")
     output   = os.path.join(tmp_dir, f"{uuid}.mp4")
 
-    # writer = cv2.VideoWriter(
-    #     temp_avi,
-    #     cv2.VideoWriter_fourcc(*"MJPG"),
-    #     FPS,
-    #     (w, h),
-    # )
-    # for i in range(n_frames):
-    #     writer.write(decode_frame(frames[i]))
-    # writer.release()
-
-    # Instead of fetching frame by frame
-    all_frames = frames[0:n_frames]  # single batched network read
-    
-    writer = cv2.VideoWriter(temp_avi, cv2.VideoWriter_fourcc(*"MJPG"), FPS, (w, h))
+    writer = cv2.VideoWriter(
+        temp_avi, cv2.VideoWriter_fourcc(*"MJPG"), FPS, (w, h)
+    )
     for i in range(n_frames):
         writer.write(decode_frame(all_frames[i]))
     writer.release()
@@ -139,6 +129,60 @@ def convert_zarr_from_hf(uuid: str, tmp_dir: str) -> str:
     os.remove(temp_avi)
     return output
 
+# def convert_zarr_from_hf(uuid: str, tmp_dir: str) -> str:
+#     """
+#     Open images.front_1 directly from HF via fsspec, convert to mp4.
+#     Returns path to the output mp4.
+#     """
+#     fs = HfFileSystem()
+#     store_path = f"hf://datasets/{EGOVERSE_REPO}/{uuid}/images.front_1"
+
+#     store = zarr.storage.FsspecStore.from_url(store_path)
+#     root  = zarr.open(store, mode="r", zarr_format=3)
+
+#     if "c" not in root and not hasattr(root, "shape"):
+#         # root IS the array when opened at images.front_1
+#         frames = root
+#     else:
+#         frames = root  # images.front_1 is the array itself
+
+#     n_frames = frames.shape[0]
+#     first    = decode_frame(frames[0])
+#     h, w     = first.shape[:2]
+
+#     temp_avi = os.path.join(tmp_dir, f"{uuid}.avi")
+#     output   = os.path.join(tmp_dir, f"{uuid}.mp4")
+
+#     # writer = cv2.VideoWriter(
+#     #     temp_avi,
+#     #     cv2.VideoWriter_fourcc(*"MJPG"),
+#     #     FPS,
+#     #     (w, h),
+#     # )
+#     # for i in range(n_frames):
+#     #     writer.write(decode_frame(frames[i]))
+#     # writer.release()
+
+#     # Instead of fetching frame by frame
+#     all_frames = frames[0:n_frames]  # single batched network read
+    
+#     writer = cv2.VideoWriter(temp_avi, cv2.VideoWriter_fourcc(*"MJPG"), FPS, (w, h))
+#     for i in range(n_frames):
+#         writer.write(decode_frame(all_frames[i]))
+#     writer.release()
+
+#     subprocess.run(
+#         ["ffmpeg", "-y",
+#          "-i", temp_avi,
+#          "-c:v", "libx264", "-preset", "fast",
+#          "-crf", "18", "-pix_fmt", "yuv420p",
+#          "-movflags", "+faststart",
+#          output],
+#         check=True, capture_output=True,
+#     )
+#     os.remove(temp_avi)
+#     return output
+
 # ─────────────────────────────────────────────
 # MCAP helpers (abc-teleop)
 # ─────────────────────────────────────────────
@@ -207,8 +251,16 @@ def stream_videos(dataset: str, uuid_input: str, task: str):
     tmp_dir = tempfile.mkdtemp(prefix="viewer_")
     try:
         # ── Egoverse (zarr → single front camera) ───────────────────────
+        # if dataset == "EgoVerse":
+        #     yield f"Fetching zarr for {uuid} ...", None, None, None
+        #     try:
+        #         mp4 = convert_zarr_from_hf(uuid, tmp_dir)
+        #         yield f"Done: {uuid}", mp4, None, None
+        #     except Exception as e:
+        #         yield f"Error: {e}", None, None, None
+
         if dataset == "EgoVerse":
-            yield f"Fetching zarr for {uuid} ...", None, None, None
+            yield f"Downloading zarr for {uuid} ...", None, None, None
             try:
                 mp4 = convert_zarr_from_hf(uuid, tmp_dir)
                 yield f"Done: {uuid}", mp4, None, None