Update app.py

app.py

@@ -7,9 +7,111 @@ import numpy as np
 import cv2
 import base64
 from scipy.ndimage import gaussian_filter1d
+import requests
+import json
+import tempfile
+import os
+from pathlib import Path
+from typing import Tuple, Optional
+from urllib.parse import urljoin
+# ------------------ 下载数据 ------------------
 
+class RemoteDatasetLoader:
+    """从 Hugging Face Hub 远程加载数据集的类"""
+
+    def __init__(self, repo_id: str, timeout: int = 30):
+        self.repo_id = repo_id
+        self.timeout = timeout
+        self.base_url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/"
+
+    def _get_dataset_info(self) -> dict:
+        info_url = urljoin(self.base_url, "meta/info.json")
+        response = requests.get(info_url, timeout=self.timeout)
+        response.raise_for_status()
+        return response.json()
+
+    def _get_episode_info(self, episode_id: int) -> dict:
+        episodes_url = urljoin(self.base_url, "meta/episodes.jsonl")
+        response = requests.get(episodes_url, timeout=self.timeout)
+        response.raise_for_status()
+        episodes = [json.loads(line) for line in response.text.splitlines() if line.strip()]
+        for episode in episodes:
+            if episode.get("episode_index") == episode_id:
+                return episode
+        raise ValueError(f"Episode {episode_id} not found")
+
+    def _download_video(self, video_url: str, save_path: str) -> str:
+        response = requests.get(video_url, timeout=self.timeout, stream=True)
+        response.raise_for_status()
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        with open(save_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        return save_path
+
+    def load_episode_data(self, episode_id: int,
+                          video_keys: Optional[list] = None,
+                          download_dir: Optional[str] = None) -> Tuple[list, pd.DataFrame]:
+        dataset_info = self._get_dataset_info()
+        episode_info = self._get_episode_info(episode_id)
+
+        if download_dir is None:
+            download_dir = tempfile.mkdtemp(prefix="lerobot_videos_")
+
+        if video_keys is None:
+            video_keys = [key for key, feature in dataset_info["features"].items()
+                          if feature["dtype"] == "video"]
+
+        video_keys = video_keys[:2]
+        video_paths = []
+        chunks_size = dataset_info.get("chunks_size", 1000)
+
+        for i, video_key in enumerate(video_keys):
+            video_url = self.base_url + dataset_info["video_path"].format(
+                episode_chunk=episode_id // chunks_size,
+                video_key=video_key,
+                episode_index=episode_id
+            )
+            video_filename = f"episode_{episode_id}_{video_key}.mp4"
+            local_path = os.path.join(download_dir, video_filename)
+            try:
+                downloaded_path = self._download_video(video_url, local_path)
+                video_paths.append(downloaded_path)
+                print(f"Downloaded video {i+1}: {downloaded_path}")
+            except Exception as e:
+                print(f"Failed to download video {video_key}: {e}")
+                video_paths.append(video_url)
+
+        data_url = self.base_url + dataset_info["data_path"].format(
+            episode_chunk=episode_id // chunks_size,
+            episode_index=episode_id
+        )
+        try:
+            df = pd.read_parquet(data_url)
+            print(f"Loaded data with {len(df)} rows and {len(df.columns)} columns")
+        except Exception as e:
+            print(f"Failed to load data: {e}")
+            df = pd.DataFrame()
+
+        return video_paths, df
+
+
+def load_remote_dataset(repo_id: str,
+                        episode_id: int = 0,
+                        video_keys: Optional[list] = None,
+                        download_dir: Optional[str] = None) -> Tuple[list, pd.DataFrame]:
+    loader = RemoteDatasetLoader(repo_id)
+    return loader.load_episode_data(episode_id, video_keys, download_dir)
+
+
+video_paths, data_df = load_remote_dataset(
+            repo_id="zijian2022/sortingtest",
+            episode_id=0,
+            download_dir="./downloaded_videos"
+        )
 # ------------------ 加载数据 ------------------
-df = pd.read_parquet("./data/clean_data/uni_boxing_object_vfm/data/chunk-000/episode_000000.parquet")
+#df = pd.read_parquet("./data/clean_data/uni_boxing_object_vfm/data/chunk-000/episode_000000.parquet")
+df = data_df
 columns = ["shoulder_pan", "shoulder_pitch", "elbow", "wrist_pitch", "wrist_roll", "gripper"]
 timestamps = df["timestamp"].values
 delta_t = np.diff(timestamps)
@@ -17,9 +119,10 @@ time_for_plot = timestamps[1:]
 action_df = pd.DataFrame(df["action"].tolist(), columns=columns)
 
 # ------------------ 视频路径 ------------------
-video_path_1 = "./data/clean_data/uni_boxing_object_vfm/videos/chunk-000/observation.images.laptop/episode_000000.mp4"
-video_path_2 = "./data/clean_data/uni_boxing_object_vfm/videos/chunk-000/observation.images.phone/episode_000000.mp4"
-
+#video_path_1 = "./data/clean_data/uni_boxing_object_vfm/videos/chunk-000/observation.images.laptop/episode_000000.mp4"
+#video_path_2 = "./data/clean_data/uni_boxing_object_vfm/videos/chunk-000/observation.images.phone/episode_000000.mp4"
+video_path_1 = video_paths[0]
+video_path_2 = video_paths[1] 
 # ------------------ Dash 初始化 ------------------
 app = dash.Dash(__name__)
 server = app.server
@@ -374,5 +477,5 @@ def update_video_frames(*args):
         return [no_update]*12 + [hover_state]
 
 # ------------------ 启动应用 ------------------
-if __name__ == '__main__':
-    app.run_server(host="0.0.0.0", port=7860, debug=False)
+if __name__ == "__main__":
+    app.run(debug=True)