push app

app.py

@@ -0,0 +1,201 @@
+import gradio as gr
+from datasets import load_dataset
+import os
+
+def load_rfm_dataset(dataset_name, config_name):
+    """Load the RFM dataset from HuggingFace Hub."""
+    try:
+        dataset = load_dataset(dataset_name, config_name, split="train")
+        return dataset, f"✅ Loaded {len(dataset)} trajectories from {dataset_name}/{config_name}"
+    except Exception as e:
+        return None, f"❌ Error loading dataset: {e}"
+
+def get_available_configs(dataset_name):
+    """Get available configurations for a dataset."""
+    try:
+        dataset_info = load_dataset(dataset_name, trust_remote_code=True)
+        configs = list(dataset_info.keys())
+        return configs
+    except Exception as e:
+        return []
+
+def visualize_trajectory(dataset, index):
+    """
+    Function to retrieve a trajectory and its metadata from the dataset.
+    """
+    if dataset is None:
+        return None, "Error: Could not load dataset", "Error: Could not load dataset"
+    
+    try:
+        item = dataset[int(index)]
+        
+        # Get the video file path
+        video_path = item["frames"]
+        
+        # Get metadata
+        task = item["task"]
+        optimal = item["optimal"]
+        is_robot = item["is_robot"]
+        data_source = item["data_source"]
+        
+        # Create metadata text
+        metadata = f"""
+        ## Trajectory Information
+        
+        **Task:** {task}
+        
+        **Optimality:** {optimal}
+        
+        **Data Type:** {'🤖 Robot' if is_robot else '👤 Human'}
+        
+        **Source:** {data_source}
+        
+        **Video Path:** `{video_path}`
+        
+        **Trajectory ID:** {item.get('id', 'N/A')}
+        """
+        
+        return video_path, metadata, f"Trajectory {index}"
+        
+    except Exception as e:
+        return None, f"Error: {str(e)}", f"Error: {str(e)}"
+
+# Create the Gradio interface
+with gr.Blocks(title="RFM Dataset Visualizer") as demo:
+    gr.Markdown("# RFM Dataset Visualizer")
+    gr.Markdown("Browse through trajectory videos and their metadata from the Reward Foundation Model dataset.")
+    
+    # Dataset selection
+    with gr.Row():
+        with gr.Column(scale=1):
+            dataset_name_input = gr.Textbox(
+                value="aliangdw/rfm",
+                label="Dataset Name",
+                placeholder="username/dataset-name"
+            )
+        
+        with gr.Column(scale=1):
+            config_name_input = gr.Textbox(
+                value="libero_10",
+                label="Configuration Name",
+                placeholder="config-name"
+            )
+        
+        with gr.Column(scale=1):
+            load_btn = gr.Button("Load Dataset", variant="primary")
+    
+    # Status message
+    status_output = gr.Markdown("Ready to load dataset...")
+    
+    # Dataset info
+    dataset_info = gr.Markdown("")
+    
+    # Visualization section
+    with gr.Row():
+        with gr.Column(scale=2):
+            # Video display
+            video_output = gr.Video(label="Trajectory Video", height=400)
+        
+        with gr.Column(scale=1):
+            # Metadata display
+            metadata_output = gr.Markdown(label="Metadata")
+    
+    # Navigation controls
+    with gr.Row():
+        with gr.Column(scale=1):
+            prev_btn = gr.Button("⬅️ Previous", variant="secondary")
+        
+        with gr.Column(scale=2):
+            # Slider for navigation
+            slider = gr.Slider(
+                minimum=0,
+                maximum=0,
+                step=1,
+                value=0,
+                label="Trajectory Index"
+            )
+        
+        with gr.Column(scale=1):
+            next_btn = gr.Button("Next ➡️", variant="secondary")
+    
+    # Current trajectory title
+    title_output = gr.Textbox(label="Current Trajectory", interactive=False)
+    
+    # State variables
+    current_dataset = gr.State(None)
+    current_index = gr.State(0)
+    
+    def load_dataset(dataset_name, config_name):
+        """Load the dataset and update the interface."""
+        dataset, status = load_rfm_dataset(dataset_name, config_name)
+        if dataset is not None:
+            max_index = len(dataset) - 1
+            info = f"**Dataset Info:**\n- **Total Trajectories:** {len(dataset)}\n- **Features:** {list(dataset.features.keys())}"
+            return dataset, status, info, max_index, 0
+        else:
+            return None, status, "", 0, 0
+    
+    def update_trajectory(dataset, index):
+        """Update the displayed trajectory."""
+        if dataset is None:
+            return None, "No dataset loaded", "No dataset loaded"
+        return visualize_trajectory(dataset, index)
+    
+    def next_trajectory(dataset, current_idx):
+        """Go to next trajectory."""
+        if dataset is None:
+            return current_idx, None, "No dataset loaded", "No dataset loaded"
+        next_idx = min(current_idx + 1, len(dataset) - 1)
+        video, metadata, title = visualize_trajectory(dataset, next_idx)
+        return next_idx, video, metadata, title
+    
+    def prev_trajectory(dataset, current_idx):
+        """Go to previous trajectory."""
+        if dataset is None:
+            return current_idx, None, "No dataset loaded", "No dataset loaded"
+        prev_idx = max(current_idx - 1, 0)
+        video, metadata, title = visualize_trajectory(dataset, prev_idx)
+        return prev_idx, video, metadata, title
+    
+    # Connect the components
+    load_btn.click(
+        fn=load_dataset,
+        inputs=[dataset_name_input, config_name_input],
+        outputs=[current_dataset, status_output, dataset_info, slider, current_index]
+    )
+    
+    slider.change(
+        fn=lambda dataset, idx: update_trajectory(dataset, idx),
+        inputs=[current_dataset, slider],
+        outputs=[video_output, metadata_output, title_output]
+    )
+    
+    next_btn.click(
+        fn=next_trajectory,
+        inputs=[current_dataset, current_index],
+        outputs=[current_index, video_output, metadata_output, title_output]
+    ).then(
+        fn=lambda idx: idx,
+        inputs=current_index,
+        outputs=slider
+    )
+    
+    prev_btn.click(
+        fn=prev_trajectory,
+        inputs=[current_dataset, current_index],
+        outputs=[current_index, video_output, metadata_output, title_output]
+    ).then(
+        fn=lambda idx: idx,
+        inputs=current_index,
+        outputs=slider
+    )
+    
+    # Load initial dataset
+    demo.load(
+        fn=lambda: load_dataset("aliangdw/rfm", "libero_10"),
+        outputs=[current_dataset, status_output, dataset_info, slider, current_index]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()