robometer/rewardeval_ui
4
1#!/usr/bin/env python32from __future__ import annotations3 4import torch5import io6import json7import os8from pathlib import Path9from typing import Any, Dict, List, Union, Optional, Tuple10from datetime import datetime11 12import aiohttp13import numpy as np14import requests15import torch16 17from dataset_types import PreferenceSample, ProgressSample, Trajectory18 19 20def pad_trajectory_to_max_frames_np(21 frames: np.ndarray, progress: List[float], max_frames: int, pad_from: str = "right"22) -> Tuple[np.ndarray, List[float]]:23 """Pad trajectory frames and progress to max_frames by repeating the first frame/progress if needed.24 25 Args:26 frames: Trajectory frames (numpy array)27 progress: Progress values (list of floats)28 max_frames: Target number of frames29 30 Returns:31 Tuple[np.ndarray, List[float]: (padded_frames, padded_progress)32 """33 current_frames = frames.shape[0]34 35 if current_frames >= max_frames:36 # No padding needed37 return frames, progress38 39 if pad_from == "left":40 pad_frame = frames[0:1] # Keep the batch dimension41 pad_progress = progress[0]42 else:43 pad_frame = frames[-1:]44 pad_progress = progress[-1]45 46 # Calculate how many frames to pad47 frames_to_pad = max_frames - current_frames48 49 # Pad frames by repeating the first frame50 if pad_from == "left":51 padded_frames = np.concatenate([np.repeat(pad_frame, frames_to_pad, axis=0), frames], axis=0)52 padded_progress = [pad_progress] * frames_to_pad + progress53 else:54 padded_frames = np.concatenate([frames, np.repeat(pad_frame, frames_to_pad, axis=0)], axis=0)55 padded_progress = progress + [pad_progress] * frames_to_pad56 57 return padded_frames, padded_progress58 59 60def linspace_subsample_frames(61 frames: np.ndarray, num_frames: int = 8, end_idx: Optional[int] = None62) -> Tuple[np.ndarray, List[int]]:63 """Uniformly subsample frames from a trajectory and return the indices.64 65 This method takes the full trajectory (e.g., 64 frames) and uniformly subsamples66 num_frames from it. The first and last frames are always included.67 68 Args:69 frames: Full trajectory frames (N frames)70 num_frames: Number of frames to subsample (default: 8)71 end_idx: Optional end index to subsample up to (if None, uses total_frames - 1)72 73 Returns:74 Tuple[np.ndarray, List[int]: (subsampled_frames, subsampled_indices)75 """76 if hasattr(frames, "shape"):77 total_frames = frames.shape[0]78 else:79 total_frames = len(frames)80 81 if total_frames <= 0:82 return frames, []83 84 # Use end_idx if provided, otherwise use full trajectory85 if end_idx is not None:86 end_idx = min(end_idx, total_frames - 1)87 frames_to_subsample = frames[: end_idx + 1]88 effective_total = end_idx + 189 else:90 frames_to_subsample = frames91 effective_total = total_frames92 93 if effective_total <= num_frames:94 # If we have fewer (or equal) frames than requested, return all frames95 indices = list(range(effective_total))96 return frames_to_subsample, indices97 98 # Special case: if num_frames == 1, always take the last frame99 if num_frames == 1:100 indices = [effective_total - 1]101 subsampled_frames = frames_to_subsample[indices]102 return subsampled_frames, indices103 104 # Evenly spaced indices from 0 to effective_total-1, inclusive105 indices_np = np.linspace(0, effective_total - 1, num_frames)106 indices = np.rint(indices_np).astype(int).tolist()107 108 # Enforce first and last explicitly109 indices[0] = 0110 indices[-1] = effective_total - 1111 112 # Ensure indices are strictly non-decreasing and within bounds113 for k in range(1, len(indices)):114 if indices[k] < indices[k - 1]:115 indices[k] = indices[k - 1]116 if indices[k] >= effective_total:117 indices[k] = effective_total - 1118 119 # Subsample frames120 subsampled_frames = frames_to_subsample[indices]121 122 return subsampled_frames, indices123 124 125def raw_dict_to_sample(126 raw_data: Union[Tuple[Dict[str, Any], Dict[str, Any]], Dict[str, Any]],127 max_frames: int = 16,128 sample_type: str = "progress",129) -> Union[ProgressSample, PreferenceSample]:130 """131 Convert raw data dictionary to a ProgressSample or PreferenceSample.132 133 Args:134 raw_data: Dict with 'frames', 'task', 'id', 'metadata', 'video_embeddings', 'text_embedding' or Tuple of (Dict[str, Any], Dict[str, Any])135 max_frames: Maximum number of frames to use (default: 16)136 sample_type: Either "progress" or "preference" (default: "progress")137 138 Returns:139 ProgressSample or PreferenceSample140 """141 142 def _build_trajectory(raw_data: Dict[str, Any], num_frames: int) -> Trajectory:143 processed_item: Dict[str, Any] = {}144 145 # Process frames146 frames_array = raw_data["frames"]147 148 # Ensure we have the correct shape: (T, H, W, C)149 if len(frames_array.shape) != 4:150 raise ValueError(f"Expected 4D array (T, H, W, C), got shape {frames_array.shape}")151 152 # Convert from CxHxW to HxWxC if needed153 if frames_array.shape[1] == 3:154 frames_array = np.transpose(frames_array, (0, 2, 3, 1))155 156 frames_array, _ = linspace_subsample_frames(frames_array, num_frames)157 dummy_progress = [0.0] * len(frames_array)158 frames_array, _ = pad_trajectory_to_max_frames_np(frames_array, dummy_progress, num_frames, pad_from="right")159 160 if frames_array.size == 0:161 raise ValueError("No frames processed for example")162 163 processed_item["frames"] = frames_array164 processed_item["frames_shape"] = frames_array.shape165 processed_item["task"] = raw_data["task"]166 processed_item["lang_vector"] = None167 processed_item["metadata"] = raw_data.get("metadata", None)168 169 # Process video embeddings using same helper functions170 video_embeddings = raw_data.get("video_embeddings")171 if video_embeddings is not None:172 video_embeddings, _ = linspace_subsample_frames(video_embeddings, num_frames)173 dummy_progress_emb = [0.0] * len(video_embeddings)174 video_embeddings, _ = pad_trajectory_to_max_frames_np(175 video_embeddings, dummy_progress_emb, num_frames, pad_from="right"176 )177 178 text_embedding = raw_data.get("text_embedding")179 180 # Convert to tensors if they are numpy arrays181 if video_embeddings is not None and isinstance(video_embeddings, np.ndarray):182 video_embeddings = torch.tensor(video_embeddings)183 if text_embedding is not None and isinstance(text_embedding, np.ndarray):184 text_embedding = torch.tensor(text_embedding)185 186 processed_item["video_embeddings"] = video_embeddings187 processed_item["text_embedding"] = text_embedding188 processed_item["video_shape"] = video_embeddings.shape if video_embeddings is not None else None189 processed_item["text_shape"] = text_embedding.shape if text_embedding is not None else None190 191 trajectory = Trajectory(**processed_item)192 return trajectory193 194 if sample_type == "progress":195 assert isinstance(raw_data, dict), "raw_data must be a dictionary"196 trajectory = _build_trajectory(raw_data=raw_data, num_frames=max_frames)197 return ProgressSample(trajectory=trajectory)198 elif sample_type == "preference":199 assert isinstance(raw_data, tuple), "raw_data must be a tuple"200 assert len(raw_data) == 2, "raw_data must be a tuple of two dictionaries"201 trajectories: List[Trajectory] = []202 for trajectory_data in raw_data:203 trajectory = _build_trajectory(raw_data=trajectory_data, num_frames=max_frames)204 trajectories.append(trajectory)205 return PreferenceSample(chosen_trajectory=trajectories[0], rejected_trajectory=trajectories[1])206 else:207 raise ValueError(f"Unsupported sample_type: {sample_type}")208 209 210def build_payload(211 samples: list[PreferenceSample | ProgressSample],212) -> tuple[dict[str, Any], list[dict[str, Any]]]:213 """Build a payload with numpy array handling.214 215 Args:216 samples: List of samples to convert217 218 Returns:219 Tuple of (files, sample_data) where:220 - files: Dict of numpy arrays converted to .npy format221 - sample_data: List of sample dictionaries with numpy arrays replaced by file references222 """223 files = {}224 sample_data = []225 226 for sample_idx, sample in enumerate(samples):227 # Copy the original sample and handle numpy arrays228 processed_sample = sample.model_dump().copy()229 230 # Handle trajectory objects with numpy arrays231 for key in [232 "chosen_trajectory",233 "rejected_trajectory",234 "trajectory",235 ]:236 if key in processed_sample and isinstance(processed_sample[key], dict):237 trajectory = processed_sample[key]238 239 # Convert numpy arrays to .npy files240 numpy_fields = ["frames", "lang_vector", "video_embeddings", "text_embedding"]241 for field_name in numpy_fields:242 # if it is a tensor, first convert it to a numpy array243 if field_name in trajectory and isinstance(trajectory[field_name], torch.Tensor):244 trajectory[field_name] = trajectory[field_name].numpy()245 246 if field_name in trajectory and isinstance(trajectory[field_name], np.ndarray):247 # Convert numpy array to .npy file248 buf = io.BytesIO()249 np.save(buf, trajectory[field_name])250 buf.seek(0)251 file_key = f"sample_{sample_idx}_{key}_{field_name}"252 files[file_key] = (253 f"sample_{sample_idx}_{key}_{field_name}.npy",254 buf,255 "application/octet-stream",256 )257 trajectory[field_name] = {"__numpy_file__": file_key}258 259 sample_data.append(processed_sample)260 261 return files, sample_data262 263 264def post_batch(url: str, payload: dict[str, Any], timeout_s: float = 120.0) -> dict[str, Any]:265 """POST a batch payload to the evaluation server and return parsed JSON."""266 resp = requests.post(url.rstrip("/") + "/evaluate_batch", json=payload, timeout=timeout_s)267 resp.raise_for_status()268 return resp.json()269 270 271def post_batch_npy(272 url: str,273 files: dict[str, Any],274 sample_data: list[dict[str, Any]],275 timeout_s: float = 120.0,276 extra_form_data: Optional[dict[str, Any]] = None,277) -> dict[str, Any]:278 """POST batch using .npy format for numpy arrays.279 280 Args:281 url: Server URL282 files: Dict of numpy arrays converted to .npy format283 sample_data: List of sample dictionaries284 timeout_s: Request timeout in seconds285 extra_form_data: Optional extra form data to include (e.g., use_frame_steps)286 """287 # Convert sample_data to form data288 data = {f"sample_{i}": json.dumps(sample) for i, sample in enumerate(sample_data)}289 290 # Add extra form data if provided291 if extra_form_data:292 for key, value in extra_form_data.items():293 data[key] = json.dumps(value) if not isinstance(value, str) else value294 295 # Send as multipart form data296 resp = requests.post(url.rstrip("/") + "/evaluate_batch_npy", files=files, data=data, timeout=timeout_s)297 resp.raise_for_status()298 return resp.json()299 300 301async def post_batch_npy_async(302 session: aiohttp.ClientSession,303 url: str,304 files: dict[str, Any],305 sample_data: list[dict[str, Any]],306 timeout_s: float = 120.0,307) -> dict[str, Any]:308 """Async version of post_batch_npy using aiohttp."""309 # Create FormData for aiohttp310 form_data = aiohttp.FormData()311 312 # Add files313 for key, (filename, file_obj, content_type) in files.items():314 form_data.add_field(key, file_obj, filename=filename, content_type=content_type)315 316 # Add sample data317 for i, sample in enumerate(sample_data):318 form_data.add_field(f"sample_{i}", json.dumps(sample))319 320 headers = {"Connection": "close"}321 # Send as multipart form data using aiohttp322 timeout = aiohttp.ClientTimeout(total=timeout_s)323 async with session.post(324 url.rstrip("/") + "/evaluate_batch_npy", data=form_data, timeout=timeout, headers=headers325 ) as resp:326 resp.raise_for_status()327 return await resp.json()328 329 330async def parse_npy_form_data(form_data: Any) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:331 """Parse multipart form data to extract numpy arrays and other data.332 333 Args:334 form_data: FastAPI form data from request.form()335 336 Returns:337 Tuple of (numpy_arrays dict, other_data dict)338 """339 numpy_arrays = {}340 other_data = {}341 342 for key, value in form_data.items():343 # Check if this is a file upload (UploadFile object)344 if hasattr(value, "filename") and value.filename:345 # This is a file upload346 if value.filename.endswith(".npy"):347 # Load .npy file (await async read)348 content = await value.read()349 buf = io.BytesIO(content)350 array = np.load(buf)351 numpy_arrays[key] = array352 else:353 # Non-.npy file, skip for now354 continue355 else:356 # This is a string value (form field)357 try:358 # Try to parse as JSON359 other_data[key] = json.loads(value)360 except (json.JSONDecodeError, TypeError):361 # Keep as string if not JSON362 other_data[key] = value363 364 return numpy_arrays, other_data365 366 367def reconstruct_payload_from_npy(368 numpy_arrays: Dict[str, np.ndarray],369 other_data: Dict[str, Any],370 trajectory_keys: Optional[List[str]] = None,371 convert_embeddings_to_torch: bool = False,372) -> List[Dict[str, Any]]:373 """Reconstruct the original payload structure from .npy files and form data.374 375 The client sends data in this format:376 - Files: sample_0_chosen_trajectory_frames.npy, sample_0_trajectory_frames.npy, etc.377 - Data: sample_0, sample_1, etc. (each containing the full sample JSON with numpy file references)378 379 Args:380 numpy_arrays: Dictionary of numpy arrays loaded from .npy files381 other_data: Dictionary of other form data382 trajectory_keys: List of trajectory keys to process (default: common keys)383 convert_embeddings_to_torch: Whether to convert embeddings to torch tensors384 385 Returns:386 List of reconstructed sample dictionaries387 """388 if trajectory_keys is None:389 trajectory_keys = [390 "chosen_trajectory",391 "rejected_trajectory",392 "trajectory",393 ]394 395 samples = []396 397 # Process each sample398 for i in range(len(other_data)):399 sample_key = f"sample_{i}"400 if sample_key in other_data:401 # Get the sample data - might already be parsed or might be a string402 sample_data = other_data[sample_key]403 if isinstance(sample_data, str):404 # Parse the sample JSON if it's a string405 sample_data = json.loads(sample_data)406 407 # Replace numpy file references with actual arrays408 for key, value in sample_data.items():409 if key in trajectory_keys:410 if isinstance(value, dict):411 for traj_key, traj_value in value.items():412 if isinstance(traj_value, dict) and traj_value.get("__numpy_file__"):413 # Replace with actual numpy array414 file_key = traj_value["__numpy_file__"]415 if file_key in numpy_arrays:416 value[traj_key] = numpy_arrays[file_key]417 418 # Convert embeddings to torch if requested419 if convert_embeddings_to_torch and traj_key in ["video_embeddings", "text_embedding"]:420 if traj_key in value and value[traj_key] is not None:421 if isinstance(value[traj_key], np.ndarray):422 value[traj_key] = torch.tensor(value[traj_key])423 elif isinstance(value[traj_key], list):424 value[traj_key] = torch.tensor(value[traj_key])425 426 samples.append(sample_data)427 428 return samples429 430 431def find_video_files(directory: str) -> list[str]:432 """Find all video files in a directory.433 434 Args:435 directory: Path to directory containing video files436 437 Returns:438 List of paths to video files439 """440 video_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv", ".wmv", ".m4v"}441 video_files = []442 443 directory_path = Path(directory)444 if not directory_path.is_dir():445 return []446 447 for file_path in directory_path.iterdir():448 if file_path.is_file() and file_path.suffix.lower() in video_extensions:449 video_files.append(str(file_path))450 451 video_files.sort()452 return video_files453 454 455def infer_task_from_video_name(video_path: str) -> str:456 """Infer task name from video filename.457 458 Task is everything before the comma (if comma exists), or everything before success/fail/failure.459 460 Args:461 video_path: Path to video file462 463 Returns:464 Inferred task name465 """466 video_name = Path(video_path).stem # Get filename without extension467 468 # If there's a comma, task is everything before the comma469 if "," in video_name:470 task_part = video_name.split(",")[0]471 else:472 # Otherwise, split by underscore and remove success/fail/failure suffixes473 parts = video_name.split("_")474 filtered_parts = []475 for part in parts:476 part_lower = part.lower()477 if part_lower not in ["success", "fail", "failure"]:478 filtered_parts.append(part)479 480 if not filtered_parts:481 return "Complete the task"482 483 task_part = "_".join(filtered_parts)484 485 # Split by underscore and join with spaces486 task_words = task_part.split("_")487 task = " ".join(task_words)488 489 if task:490 # Capitalize first letter of first word, keep rest as is491 task = task[0].upper() + task[1:] if len(task) > 1 else task.upper()492 else:493 task = "Complete the task"494 495 return task496 497 498def setup_output_directory(output_dir: Optional[str], video_path: Optional[str] = None) -> str:499 """Create output directory and return path."""500 if output_dir:501 save_dir = output_dir502 else:503 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")504 save_dir = os.path.join(".", f"eval_outputs/{timestamp}")505 506 os.makedirs(save_dir, exist_ok=True)507 return save_dir508