select videos

app.py

@@ -9,6 +9,7 @@ import os
 import tempfile
 from pathlib import Path
 from typing import Optional, Tuple
+import logging
 
 import gradio as gr
 import spaces  # Required for ZeroGPU
@@ -23,16 +24,51 @@ from typing import Any, Optional, Tuple
 
 from rfm.data.dataset_types import Trajectory, ProgressSample, PreferenceSample
 from rfm.evals.eval_utils import build_payload, post_batch_npy
+from datasets import load_dataset as load_dataset_hf, get_dataset_config_names
+
+logger = logging.getLogger(__name__)
+
+# Predefined dataset names (same as visualizer)
+PREDEFINED_DATASETS = [
+    "abraranwar/agibotworld_alpha_rfm",
+    "abraranwar/libero_rfm",
+    "abraranwar/usc_koch_rewind_rfm",
+    "aliangdw/metaworld",
+    "anqil/rh20t_rfm",
+    "anqil/rh20t_subset_rfm",
+    "jesbu1/auto_eval_rfm",
+    "jesbu1/egodex_rfm",
+    "jesbu1/epic_rfm",
+    "jesbu1/fino_net_rfm",
+    "jesbu1/failsafe_rfm",
+    "jesbu1/hand_paired_rfm",
+    "jesbu1/galaxea_rfm",
+    "jesbu1/h2r_rfm",
+    "jesbu1/humanoid_everyday_rfm",
+    "jesbu1/molmoact_rfm",
+    "jesbu1/motif_rfm",
+    "jesbu1/oxe_rfm",
+    "jesbu1/oxe_rfm_eval",
+    "jesbu1/ph2d_rfm",
+    "jesbu1/racer_rfm",
+    "jesbu1/roboarena_0825_rfm",
+    "jesbu1/soar_rfm",
+    "ykorkmaz/libero_failure_rfm",
+    "aliangdw/usc_xarm_policy_ranking",
+    "aliangdw/usc_franka_policy_ranking",
+    "aliangdw/utd_so101_policy_ranking",
+    "aliangdw/utd_so101_human"
+]
 
 # Global server state
 _server_state = {
     "server_url": None,
 }
 
-def check_server_health(server_url: str) -> Tuple[str, Optional[dict]]:
+def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[str]]:
     """Check server health and get model info."""
     if not server_url:
-        return "Please provide a server URL.", None
+        return "Please provide a server URL.", None, None
     
     try:
         url = server_url.rstrip("/") + "/health"
@@ -50,24 +86,187 @@ def check_server_health(server_url: str) -> Tuple[str, Optional[dict]]:
         except:
             pass
         
+        # Try to get model info
+        model_info_text = None
+        try:
+            model_info_url = server_url.rstrip("/") + "/model_info"
+            model_info_response = requests.get(model_info_url, timeout=5.0)
+            if model_info_response.status_code == 200:
+                model_info_data = model_info_response.json()
+                model_info_text = format_model_info(model_info_data)
+        except Exception as e:
+            logger.warning(f"Could not fetch model info: {e}")
+        
         _server_state["server_url"] = server_url
-        return f"Server connected: {health_data.get('available_gpus', 0)}/{health_data.get('total_gpus', 0)} GPUs available", health_data
+        return f"Server connected: {health_data.get('available_gpus', 0)}/{health_data.get('total_gpus', 0)} GPUs available", health_data, model_info_text
     except requests.exceptions.RequestException as e:
-        return f"Error connecting to server: {str(e)}", None
+        return f"Error connecting to server: {str(e)}", None, None
+
+
+def format_model_info(model_info: dict) -> str:
+    """Format model info and experiment config as markdown."""
+    lines = ["## Model Information\n"]
+    
+    # Model path
+    model_path = model_info.get("model_path", "Unknown")
+    lines.append(f"**Model Path:** `{model_path}`\n")
+    
+    # Number of GPUs
+    num_gpus = model_info.get("num_gpus", "Unknown")
+    lines.append(f"**Number of GPUs:** {num_gpus}\n")
+    
+    # Model architecture
+    model_arch = model_info.get("model_architecture", {})
+    if model_arch and "error" not in model_arch:
+        lines.append("\n## Model Architecture\n")
+        
+        model_class = model_arch.get("model_class", "Unknown")
+        model_module = model_arch.get("model_module", "Unknown")
+        lines.append(f"- **Model Class:** `{model_class}`\n")
+        lines.append(f"- **Module:** `{model_module}`\n")
+        
+        # Parameter counts
+        total_params = model_arch.get("total_parameters")
+        trainable_params = model_arch.get("trainable_parameters")
+        frozen_params = model_arch.get("frozen_parameters")
+        trainable_pct = model_arch.get("trainable_percentage")
+        
+        if total_params is not None:
+            lines.append(f"\n### Parameter Statistics\n")
+            lines.append(f"- **Total Parameters:** {total_params:,}\n")
+            if trainable_params is not None:
+                lines.append(f"- **Trainable Parameters:** {trainable_params:,}\n")
+            if frozen_params is not None:
+                lines.append(f"- **Frozen Parameters:** {frozen_params:,}\n")
+            if trainable_pct is not None:
+                lines.append(f"- **Trainable Percentage:** {trainable_pct:.2f}%\n")
+        
+        # Architecture summary
+        arch_summary = model_arch.get("architecture_summary", [])
+        if arch_summary:
+            lines.append(f"\n### Architecture Summary (Top-Level Modules)\n")
+            for module_info in arch_summary[:10]:  # Show first 10 modules
+                name = module_info.get("name", "Unknown")
+                module_type = module_info.get("type", "Unknown")
+                params = module_info.get("parameters", 0)
+                lines.append(f"- **{name}** (`{module_type}`): {params:,} parameters\n")
+    
+    # Experiment config
+    exp_config = model_info.get("experiment_config", {})
+    if exp_config:
+        lines.append("\n## Experiment Configuration\n")
+        
+        # Model config
+        model_cfg = exp_config.get("model", {})
+        if model_cfg:
+            lines.append("### Model Configuration\n")
+            lines.append(f"- **Base Model:** `{model_cfg.get('base_model_id', 'N/A')}`\n")
+            lines.append(f"- **Model Type:** `{model_cfg.get('model_type', 'N/A')}`\n")
+            lines.append(f"- **Train Progress Head:** {model_cfg.get('train_progress_head', False)}\n")
+            lines.append(f"- **Train Preference Head:** {model_cfg.get('train_preference_head', False)}\n")
+            lines.append(f"- **Train Similarity Head:** {model_cfg.get('train_similarity_head', False)}\n")
+            lines.append(f"- **Train Success Head:** {model_cfg.get('train_success_head', False)}\n")
+            lines.append(f"- **Use PEFT:** {model_cfg.get('use_peft', False)}\n")
+            lines.append(f"- **Use Unsloth:** {model_cfg.get('use_unsloth', False)}\n")
+        
+        # Data config
+        data_cfg = exp_config.get("data", {})
+        if data_cfg:
+            lines.append("\n### Data Configuration\n")
+            lines.append(f"- **Max Frames:** {data_cfg.get('max_frames', 'N/A')}\n")
+            lines.append(f"- **Resized Dimensions:** {data_cfg.get('resized_height', 'N/A')}x{data_cfg.get('resized_width', 'N/A')}\n")
+            train_datasets = data_cfg.get('train_datasets', [])
+            if train_datasets:
+                lines.append(f"- **Train Datasets:** {', '.join(train_datasets)}\n")
+            eval_datasets = data_cfg.get('eval_datasets', [])
+            if eval_datasets:
+                lines.append(f"- **Eval Datasets:** {', '.join(eval_datasets)}\n")
+        
+        # Training config
+        training_cfg = exp_config.get("training", {})
+        if training_cfg:
+            lines.append("\n### Training Configuration\n")
+            lines.append(f"- **Learning Rate:** {training_cfg.get('learning_rate', 'N/A')}\n")
+            lines.append(f"- **Batch Size:** {training_cfg.get('per_device_train_batch_size', 'N/A')}\n")
+            lines.append(f"- **Gradient Accumulation Steps:** {training_cfg.get('gradient_accumulation_steps', 'N/A')}\n")
+            lines.append(f"- **Max Steps:** {training_cfg.get('max_steps', 'N/A')}\n")
+    
+    return "".join(lines)
+
+
+def load_rfm_dataset(dataset_name, config_name):
+    """Load the RFM dataset from HuggingFace Hub."""
+    try:
+        if not dataset_name or not config_name:
+            return None, "Please provide both dataset name and configuration"
+        
+        dataset = load_dataset_hf(dataset_name, name=config_name, split="train")
+        
+        if len(dataset) == 0:
+            return None, f"Dataset {dataset_name}/{config_name} is empty"
+        
+        return dataset, f"Loaded {len(dataset)} trajectories from {dataset_name}/{config_name}"
+    except Exception as e:
+        error_msg = str(e)
+        if "not found" in error_msg.lower():
+            return None, f"Dataset or configuration not found: {dataset_name}/{config_name}"
+        elif "authentication" in error_msg.lower():
+            return None, f"Authentication required for {dataset_name}"
+        else:
+            return None, f"Error loading dataset: {error_msg}"
+
+
+def get_available_configs(dataset_name):
+    """Get available configurations for a dataset."""
+    try:
+        configs = get_dataset_config_names(dataset_name)
+        return configs
+    except Exception as e:
+        logger.warning(f"Error getting configs for {dataset_name}: {e}")
+        return []
+
+
+def get_trajectory_video_path(dataset, index, dataset_name):
+    """Get video path from a trajectory in the dataset."""
+    try:
+        item = dataset[int(index)]
+        frames_data = item["frames"]
+        
+        if isinstance(frames_data, str):
+            # Construct HuggingFace Hub URL
+            if dataset_name:
+                video_path = f"https://huggingface.co/datasets/{dataset_name}/resolve/main/{frames_data}"
+            else:
+                video_path = f"https://huggingface.co/datasets/aliangdw/rfm/resolve/main/{frames_data}"
+            return video_path, item.get("task", "Complete the task")
+        else:
+            return None, None
+    except Exception as e:
+        logger.error(f"Error getting trajectory video path: {e}")
+        return None, None
 
 
 def extract_frames(video_path: str, max_frames: int = 16, fps: float = 1.0) -> np.ndarray:
-    """Extract frames from video file as numpy array (T, H, W, C)."""
+    """Extract frames from video file as numpy array (T, H, W, C).
+    
+    Supports both local file paths and URLs (e.g., HuggingFace Hub URLs).
+    """
     if video_path is None:
         return None
 
     if isinstance(video_path, tuple):
         video_path = video_path[0]
 
-    if not os.path.exists(video_path):
+    # Check if it's a URL or local file
+    is_url = video_path.startswith(("http://", "https://"))
+    is_local_file = os.path.exists(video_path) if not is_url else False
+    
+    if not is_url and not is_local_file:
+        logger.warning(f"Video path does not exist: {video_path}")
         return None
 
     try:
+        # decord.VideoReader can handle both local files and URLs
         vr = decord.VideoReader(video_path, num_threads=1)
         total_frames = len(vr)
 
@@ -83,7 +282,7 @@ def extract_frames(video_path: str, max_frames: int = 16, fps: float = 1.0) -> n
         del vr 
         return frames_array
     except Exception as e:
-        print(f"Error extracting frames: {e}")
+        logger.error(f"Error extracting frames from {video_path}: {e}")
         return None
 
 
@@ -410,22 +609,60 @@ with demo:
                 check_connection_btn = gr.Button("Check Connection", variant="primary", size="sm")
         
         server_status = gr.Markdown("Enter server URL and click 'Check Connection'")
+        model_info_display = gr.Markdown("", visible=False)
 
         def on_check_connection(server_url: str):
             """Handle server connection check."""
-            status, health_data = check_server_health(server_url)
-            return status
+            status, health_data, model_info_text = check_server_health(server_url)
+            if model_info_text:
+                return status, gr.update(value=model_info_text, visible=True)
+            else:
+                return status, gr.update(visible=False)
         
         check_connection_btn.click(
             fn=on_check_connection,
             inputs=[server_url_input],
-            outputs=[server_status],
+            outputs=[server_status, model_info_display],
         )
 
     with gr.Tab("Progress Prediction"):
         gr.Markdown("### Progress & Success Prediction")
+        gr.Markdown("Upload a video or select one from a dataset to get progress predictions.")
+        
         with gr.Row():
             with gr.Column():
+                with gr.Accordion("📁 Select from Dataset", open=False):
+                    dataset_name_single = gr.Dropdown(
+                        choices=PREDEFINED_DATASETS,
+                        value="jesbu1/oxe_rfm",
+                        label="Dataset Name",
+                        allow_custom_value=True
+                    )
+                    config_name_single = gr.Dropdown(
+                        choices=[],
+                        value="",
+                        label="Configuration Name",
+                        allow_custom_value=True
+                    )
+                    with gr.Row():
+                        refresh_configs_btn = gr.Button("🔄 Refresh Configs", variant="secondary", size="sm")
+                        load_dataset_btn = gr.Button("Load Dataset", variant="secondary", size="sm")
+                    
+                    dataset_status_single = gr.Markdown("", visible=False)
+                    trajectory_slider = gr.Slider(
+                        minimum=0,
+                        maximum=0,
+                        step=1,
+                        value=0,
+                        label="Trajectory Index",
+                        interactive=False
+                    )
+                    use_dataset_video_btn = gr.Button("Use Selected Video", variant="secondary")
+                
+                gr.Markdown("---")
+                gr.Markdown("**OR**")
+                gr.Markdown("---")
+                
                 single_video_input = gr.Video(label="Upload Video", height=300)
                 task_text_input = gr.Textbox(
                     label="Task Description",
@@ -446,7 +683,73 @@ with demo:
                 progress_plot = gr.Image(label="Progress Prediction", height=400)
                 success_plot = gr.Image(label="Success Prediction", height=400)
                 info_output = gr.Markdown("")
-
+        
+        # State variables for dataset
+        current_dataset_single = gr.State(None)
+        
+        def update_config_choices_single(dataset_name):
+            """Update config choices when dataset changes."""
+            if not dataset_name:
+                return gr.update(choices=[], value="")
+            try:
+                configs = get_available_configs(dataset_name)
+                if configs:
+                    return gr.update(choices=configs, value=configs[0])
+                else:
+                    return gr.update(choices=[], value="")
+            except Exception as e:
+                logger.warning(f"Could not fetch configs: {e}")
+                return gr.update(choices=[], value="")
+        
+        def load_dataset_single(dataset_name, config_name):
+            """Load dataset and update slider."""
+            dataset, status = load_rfm_dataset(dataset_name, config_name)
+            if dataset is not None:
+                max_index = len(dataset) - 1
+                return (
+                    dataset,
+                    gr.update(value=status, visible=True),
+                    gr.update(maximum=max_index, value=0, interactive=True, label=f"Trajectory Index (0 to {max_index})")
+                )
+            else:
+                return None, gr.update(value=status, visible=True), gr.update(maximum=0, value=0, interactive=False)
+        
+        def use_dataset_video(dataset, index, dataset_name):
+            """Load video from dataset and update inputs."""
+            if dataset is None:
+                return None, "Complete the task", gr.update(value="No dataset loaded", visible=True)
+            
+            video_path, task = get_trajectory_video_path(dataset, index, dataset_name)
+            if video_path:
+                return video_path, task, gr.update(value=f"✅ Loaded trajectory {index} from dataset", visible=True)
+            else:
+                return None, "Complete the task", gr.update(value="❌ Error loading trajectory", visible=True)
+        
+        # Dataset selection handlers
+        dataset_name_single.change(
+            fn=update_config_choices_single,
+            inputs=[dataset_name_single],
+            outputs=[config_name_single]
+        )
+        
+        refresh_configs_btn.click(
+            fn=update_config_choices_single,
+            inputs=[dataset_name_single],
+            outputs=[config_name_single]
+        )
+        
+        load_dataset_btn.click(
+            fn=load_dataset_single,
+            inputs=[dataset_name_single, config_name_single],
+            outputs=[current_dataset_single, dataset_status_single, trajectory_slider]
+        )
+        
+        use_dataset_video_btn.click(
+            fn=use_dataset_video,
+            inputs=[current_dataset_single, trajectory_slider, dataset_name_single],
+            outputs=[single_video_input, task_text_input, dataset_status_single]
+        )
+        
         analyze_single_btn.click(
             fn=process_single_video,
             inputs=[single_video_input, task_text_input, server_url_input, fps_input_single],