import os import json from collections import defaultdict from tqdm import tqdm import numpy as np from dataset_upload.helpers import generate_unique_id trajectory_info_template = { "id": [], "task": [], # "lang_vector": [], "data_source": None, "frames": None, "is_robot": None, "quality_label": None, "partial_success": None, # in [0, 1] } class EgoCOTFrameloader: """Pickle-able loader that reads EgoCoT frames from disk on demand. Stores only simple fields so it can be safely passed across processes. """ def __init__(self, frames_path: str) -> None: self.frames_path = frames_path def __call__(self) -> np.ndarray: """Load frames from disk when called. Returns: np.ndarray of shape (T, H, W, 3), dtype uint8 """ # Load the numpy array containing 8 consecutive frames frames = np.load(self.frames_path) # Ensure the frames are in the correct format # EgoCoT frames are stored as numpy arrays with 8 consecutive frames assert frames.ndim == 4, f"Expected 4D array, got {frames.ndim}D array" # Ensure shape and dtype sanity if not isinstance(frames, np.ndarray) or frames.ndim != 4 or frames.shape[1] != 3: raise ValueError(f"Unexpected frames shape for {self.frames_path}: {getattr(frames, 'shape', None)}") # Ensure uint8 if frames.dtype != np.uint8: # Convert from float to uint8 if necessary if frames.dtype in [np.float32, np.float64]: if frames.max() <= 1.0: # Values in [0, 1] range - just scale to [0, 255] frames = (frames * 255).astype(np.uint8) else: # Values appear to be ImageNet normalized - denormalize first # ImageNet normalization: (pixel - mean) / std # Reverse: pixel * std + mean imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1) imagenet_std = np.array([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1) # Denormalize: multiply by std and add mean frames = frames * imagenet_std + imagenet_mean # Clamp to [0, 1] and convert to [0, 255] frames = np.clip(frames, 0, 1) frames = (frames * 255).astype(np.uint8) else: frames = frames.astype(np.uint8, copy=False) # now convert frames to (T, H, W, 3) from (T, C, H, W) frames = frames.transpose(0, 2, 3, 1) return frames def create_new_trajectory(frames_path: str, caption: str) -> dict: """Create a new trajectory from EgoCoT data.""" trajectory_info = {} trajectory_info["id"] = generate_unique_id() trajectory_info["task"] = caption # Use caption as the task description trajectory_info["frames"] = EgoCOTFrameloader(frames_path) trajectory_info["is_robot"] = False # EgoCoT is human egocentric data trajectory_info["quality_label"] = "successful" trajectory_info["partial_success"] = 1 trajectory_info["data_source"] = "egocot" return trajectory_info def load_egocot_dataset(dataset_path: str) -> dict[str, list[dict]]: """Load EgoCoT dataset from results.json and .npy frame files (EgoCOT_clear). Expected layout (example): / results.json EgoCOT_clear/ EGO_0000.npy EGO_0001.npy Args: dataset_path: Path to a directory containing one or more EgoCoT result folders Returns: Dictionary mapping task descriptions to lists of trajectory dictionaries """ # Locate results.json files json_files = [] for root, dirs, files in os.walk(dataset_path): for file in files: if file.lower() == "results.json": json_files.append(os.path.join(root, file)) if not json_files: raise FileNotFoundError(f"No results.json files found in {dataset_path}") task_data = defaultdict(list) total_trajectories = 0 for json_file in json_files: print(f"Loading annotations from {json_file}") with open(json_file, "r") as f: annotations = json.load(f) # Normalize JSON structures if isinstance(annotations, list): data_items = annotations elif isinstance(annotations, dict): if "data" in annotations: data_items = annotations["data"] elif "results" in annotations: data_items = annotations["results"] elif "annotations" in annotations: data_items = annotations["annotations"] elif "samples" in annotations: data_items = annotations["samples"] elif "image" in annotations: data_items = [annotations] else: raise ValueError(f"Unexpected JSON structure in {json_file}: cannot find data list") else: raise ValueError(f"Unexpected JSON structure in {json_file}") for item in tqdm(data_items, desc="Processing trajectories"): # Extract required fields (handle common variants and typos) image_filename = item.get("image") # this caption is the llm generated one that's much more detailed from EgoCOT's processing caption = item.get("planing").split("\n")[0][1:] if not caption: caption = item.get("caption") # use the backup original caption score = item.get("score") if not image_filename or not caption: print(f"Skipping item with missing image or caption: {item}") continue # Construct full path to the .npy file, prioritizing EgoCOT_clear next to results.json base_dir = os.path.dirname(json_file) candidate_paths = [ os.path.join(base_dir, "EgoCOT_clear", image_filename) if image_filename else None, # os.path.join(base_dir, image_filename) if image_filename else None, # os.path.join(dataset_path, "EgoCOT_clear", image_filename) if image_filename else None, # os.path.join(dataset_path, image_filename) if image_filename else None, ] frames_path = None for cand in candidate_paths: if cand and os.path.exists(cand): frames_path = cand break if frames_path is None: print(f"Warning: .npy frame file not found for: {image_filename}") continue if not frames_path.lower().endswith(".npy"): print(f"Warning: expected .npy file, got: {frames_path}. Skipping.") continue # Create trajectory trajectory = create_new_trajectory(frames_path, caption) # Group by task/caption for organization task_key = caption[:50] + "..." if len(caption) > 50 else caption task_data[task_key].append(trajectory) total_trajectories += 1 print(f"Loaded {total_trajectories} trajectories from {len(task_data)} unique tasks") return task_data