#!/usr/bin/env python3 """ EndoGaussian-4D Dataset Download & Organization Pipeline Downloads, validates, and organizes endoscopic datasets into a unified format. Supports: EndoNeRF, EndoSLAM, C3VD, SCARED, StereoMIS, D4D, Hamlyn. Each dataset is registered with metadata, download instructions, expected directory structure, and automatic validation. Usage: # List available datasets with access info python scripts/download_datasets.py --list # Download and organize EndoNeRF (smallest, auto-downloadable) python scripts/download_datasets.py --datasets endonerf --data-root ./data # Validate existing data installation python scripts/download_datasets.py --validate --data-root ./data # Generate manifest (JSON inventory of all available data) python scripts/download_datasets.py --manifest --data-root ./data """ import argparse import hashlib import json import os import sys import urllib.request import zipfile from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np # --------------------------------------------------------------------------- # Dataset Registry # --------------------------------------------------------------------------- @dataclass class DatasetInfo: """Metadata for a registered dataset.""" name: str description: str num_sequences: int resolution: Tuple[int, int] # (H, W) has_depth: bool has_poses: bool has_masks: bool access_method: str # "auto", "request", "registration" access_url: str download_urls: Dict[str, str] = field(default_factory=dict) expected_dirs: List[str] = field(default_factory=list) citation: str = "" notes: str = "" DATASET_REGISTRY: Dict[str, DatasetInfo] = { "endonerf": DatasetInfo( name="EndoNeRF", description="2 prostatectomy sequences (cutting/pulling) from da Vinci. " "The standard benchmark for endoscopic neural radiance fields.", num_sequences=2, resolution=(256, 320), has_depth=True, has_poses=True, has_masks=True, access_method="auto", access_url="https://github.com/med-air/EndoNeRF", download_urls={ "cutting": "https://www.dropbox.com/s/09sxbwusf0pfiks/cutting_tissues_twice.zip?dl=1", "pulling": "https://www.dropbox.com/s/i2zy2mz4lqndxco/pulling_soft_tissues.zip?dl=1", }, expected_dirs=["images", "depth", "masks"], citation="Wang et al., Neural Rendering for Stereo 3D Reconstruction of Deformable Tissues in Robotic Surgery, MICCAI 2022", notes="Includes tool masks. Direct Dropbox download available. Best for initial experiments.", ), "endoslam": DatasetInfo( name="EndoSLAM", description="76K frames from capsule/standard endoscopy with synthetic depth. " "WARNING: Depth is synthetic (Blender), not real ground truth.", num_sequences=35, resolution=(480, 640), has_depth=True, has_poses=True, has_masks=False, access_method="auto", access_url="https://github.com/CapsuleEndoscope/EndoSLAM", download_urls={}, # Need git clone expected_dirs=["Cameras", "Frames"], citation="Ozyoruk et al., EndoSLAM Dataset and An Unsupervised Monocular Visual Odometry and Depth Estimation Approach for Endoscopic Videos, Medical Image Analysis 2021", notes="Synthetic depth only. Use for pretraining/ablation, NOT for final evaluation.", ), "c3vd": DatasetInfo( name="C3VD", description="22 colonoscopy sequences on realistic phantoms with structured-light GT depth. " "High-quality ground truth from clinical-grade phantoms.", num_sequences=22, resolution=(540, 675), has_depth=True, has_poses=True, has_masks=False, access_method="request", access_url="https://durrlab.github.io/C3VD/", download_urls={}, expected_dirs=["color", "depth", "poses"], citation="Bobrow et al., Colonoscopy 3D Video Dataset with Paired Depth from 2D-3D Registration, Medical Image Analysis 2023", notes="Requires Google Form request at durrlab.github.io/C3VD/. 675×540 resolution.", ), "scared": DatasetInfo( name="SCARED", description="Surgical scene reconstruction challenge dataset from da Vinci Xi. " "Structured-light ground truth depth, requires registration.", num_sequences=7, resolution=(1024, 1280), has_depth=True, has_poses=True, has_masks=False, access_method="registration", access_url="https://endovissub2019-scared.grand-challenge.org/", download_urls={}, expected_dirs=["data", "ground_truth"], citation="Allan et al., Stereo Correspondence and Reconstruction of Endoscopic Data Challenge, 2019", notes="Requires Grand Challenge registration. High resolution stereo pairs.", ), "stereomis": DatasetInfo( name="StereoMIS", description="11 in-vivo stereo endoscopy sequences. CC-BY license on Zenodo. " "Used by Endo-4DGS and EndoGaussian as primary evaluation set.", num_sequences=11, resolution=(480, 640), has_depth=False, has_poses=False, has_masks=False, access_method="auto", access_url="https://zenodo.org/records/7727692", download_urls={}, expected_dirs=["left", "right"], citation="Hayoz et al., StereoMIS: Stereo Minimally Invasive Surgery Dataset, 2023", notes="Stereo pairs only (no GT depth). Compute depth via stereo matching. CC-BY-4.0.", ), "d4d": DatasetInfo( name="D4D (Dresden 4D)", description="98 sequences with structured-light GT depth. NEWEST dataset (March 2025). " "Fully open with DOI. Comprehensive coverage of surgical scenarios.", num_sequences=98, resolution=(480, 640), has_depth=True, has_poses=True, has_masks=True, access_method="auto", access_url="https://doi.org/10.25838/d4d-2025", download_urls={}, expected_dirs=["rgb", "depth", "segmentation"], citation="Dresden 4D Dataset, 2025", notes="March 2025 release. Fully open access. 98 sequences, structured-light GT.", ), "hamlyn": DatasetInfo( name="Hamlyn", description="Hamlyn Centre Laparoscopic/Endoscopic Video Dataset. " "Various procedures, some with stereo and tracking data.", num_sequences=20, resolution=(288, 360), has_depth=False, has_poses=False, has_masks=False, access_method="registration", access_url="http://hamlyn.doc.ic.ac.uk/vision/", download_urls={}, expected_dirs=["images"], citation="Hamlyn Centre, Imperial College London", notes="Academic use. Registration required. Lower resolution.", ), } # --------------------------------------------------------------------------- # Dataset Organizer # --------------------------------------------------------------------------- class DatasetOrganizer: """ Downloads and organizes datasets into a unified directory structure. Target structure: data_root/ ├── endonerf/ │ ├── cutting/ │ │ ├── images/ # RGB frames (PNG) │ │ ├── depth/ # Depth maps (NPY or PNG) │ │ ├── masks/ # Tool masks (PNG, binary) │ │ └── poses.npy # Camera poses [N, 3, 5] LLFF format │ └── pulling/ │ └── ... ├── c3vd/ │ ├── seq_001/ │ │ ├── images/ │ │ ├── depth/ │ │ └── poses.npy │ └── ... └── manifest.json # Inventory of all available data """ def __init__(self, data_root: str): self.data_root = Path(data_root) self.data_root.mkdir(parents=True, exist_ok=True) def list_datasets(self): """Print information about all registered datasets.""" print("\n" + "=" * 80) print("EndoGaussian-4D Dataset Registry") print("=" * 80) for key, info in DATASET_REGISTRY.items(): status = self._check_status(key) status_icon = "✓" if status == "installed" else "○" if status == "partial" else "✗" print(f"\n[{status_icon}] {info.name} ({key})") print(f" {info.description}") print(f" Sequences: {info.num_sequences} | Resolution: {info.resolution[1]}×{info.resolution[0]}") print(f" Depth: {'✓' if info.has_depth else '✗'} | " f"Poses: {'✓' if info.has_poses else '✗'} | " f"Masks: {'✓' if info.has_masks else '✗'}") print(f" Access: {info.access_method} → {info.access_url}") if info.notes: print(f" Note: {info.notes}") print("\n" + "=" * 80) print("Legend: ✓ = installed, ○ = partial, ✗ = not installed") print("Run with --datasets to download auto-accessible datasets") print("=" * 80 + "\n") def _check_status(self, dataset_key: str) -> str: """Check if a dataset is installed.""" dataset_dir = self.data_root / dataset_key if not dataset_dir.exists(): return "missing" info = DATASET_REGISTRY[dataset_key] has_any = False has_all = True for subdir in info.expected_dirs: if list(dataset_dir.rglob(subdir)): has_any = True else: has_all = False if has_all: return "installed" elif has_any: return "partial" return "missing" def download(self, dataset_key: str): """Download and organize a dataset.""" if dataset_key not in DATASET_REGISTRY: print(f"[ERROR] Unknown dataset: {dataset_key}") print(f"Available: {', '.join(DATASET_REGISTRY.keys())}") return False info = DATASET_REGISTRY[dataset_key] if info.access_method != "auto": print(f"\n[{info.name}] Cannot auto-download. Access method: {info.access_method}") print(f" Please visit: {info.access_url}") print(f" Then place files in: {self.data_root / dataset_key}") return False if not info.download_urls: print(f"\n[{info.name}] Auto-download registered but no URLs configured.") print(f" Please visit: {info.access_url}") return False dataset_dir = self.data_root / dataset_key dataset_dir.mkdir(parents=True, exist_ok=True) print(f"\n[{info.name}] Downloading {len(info.download_urls)} sequence(s)...") for seq_name, url in info.download_urls.items(): seq_dir = dataset_dir / seq_name if seq_dir.exists() and any(seq_dir.iterdir()): print(f" [{seq_name}] Already exists, skipping") continue print(f" [{seq_name}] Downloading from {url[:60]}...") zip_path = dataset_dir / f"{seq_name}.zip" try: urllib.request.urlretrieve(url, str(zip_path)) print(f" [{seq_name}] Extracting...") with zipfile.ZipFile(str(zip_path), 'r') as zf: zf.extractall(str(dataset_dir)) zip_path.unlink() print(f" [{seq_name}] Done ✓") except Exception as e: print(f" [{seq_name}] Failed: {e}") if zip_path.exists(): zip_path.unlink() return True def validate(self, dataset_key: Optional[str] = None): """Validate dataset installation.""" keys = [dataset_key] if dataset_key else list(DATASET_REGISTRY.keys()) print("\n" + "-" * 60) print("Dataset Validation Report") print("-" * 60) for key in keys: info = DATASET_REGISTRY[key] status = self._check_status(key) dataset_dir = self.data_root / key print(f"\n[{info.name}] Status: {status}") if status == "missing": print(f" Not found at {dataset_dir}") continue # Count files n_images = len(list(dataset_dir.rglob("*.png"))) + len(list(dataset_dir.rglob("*.jpg"))) n_depth = len(list(dataset_dir.rglob("*.npy"))) n_masks = 0 for mask_dir in dataset_dir.rglob("mask*"): if mask_dir.is_dir(): n_masks += len(list(mask_dir.glob("*.png"))) print(f" Images: {n_images} | Depth maps: {n_depth} | Masks: {n_masks}") print(f" Directory: {dataset_dir}") # Check for poses poses_files = list(dataset_dir.rglob("poses*.npy")) + list(dataset_dir.rglob("poses*.json")) if poses_files: print(f" Poses: {len(poses_files)} file(s) found") for pf in poses_files: if pf.suffix == ".npy": poses = np.load(str(pf)) print(f" {pf.name}: shape={poses.shape}") else: print(f" Poses: NOT FOUND — run extract_poses.py") print("\n" + "-" * 60) def generate_manifest(self) -> dict: """Generate a JSON manifest of all available data.""" manifest = { "data_root": str(self.data_root), "datasets": {}, } for key, info in DATASET_REGISTRY.items(): dataset_dir = self.data_root / key entry = { "name": info.name, "status": self._check_status(key), "has_depth": info.has_depth, "has_poses": info.has_poses, "has_masks": info.has_masks, "resolution": list(info.resolution), "sequences": [], } if dataset_dir.exists(): for seq_dir in sorted(dataset_dir.iterdir()): if seq_dir.is_dir(): n_imgs = len(list(seq_dir.rglob("*.png"))) + len(list(seq_dir.rglob("*.jpg"))) if n_imgs > 0: entry["sequences"].append({ "name": seq_dir.name, "n_frames": n_imgs, "has_poses": any(seq_dir.rglob("poses*")), }) manifest["datasets"][key] = entry manifest_path = self.data_root / "manifest.json" with open(manifest_path, "w") as f: json.dump(manifest, f, indent=2) print(f"Manifest saved to {manifest_path}") return manifest # --------------------------------------------------------------------------- # Unified EndoDataset Loader # --------------------------------------------------------------------------- @dataclass class Frame: """A single endoscopic frame with all available annotations.""" rgb: np.ndarray # [H, W, 3] uint8 depth: Optional[np.ndarray] # [H, W] float32, in mm mask: Optional[np.ndarray] # [H, W] bool (True = tissue, False = tool) pose: Optional[np.ndarray] # [4, 4] camera-to-world intrinsics: Optional[np.ndarray] # [3, 3] timestamp: float # Normalized [0, 1] frame_idx: int sequence_name: str class EndoDataset: """ Unified dataset loader supporting multiple endoscopic dataset formats. Auto-detects directory structure and loads frames with consistent API. Supported formats: - LLFF (images/ + poses_bounds.npy) - C3VD (color/ + depth/ + poses/) - EndoSLAM (Frames/ + Cameras/) - EndoNeRF (images/ + depth/ + masks/) Usage: dataset = EndoDataset("./data/endonerf/cutting") frame = dataset[0] # Frame dataclass print(frame.rgb.shape, frame.depth.shape) # Get point cloud for Gaussian initialization points, colors = dataset.get_point_cloud() """ def __init__(self, sequence_dir: str, max_frames: Optional[int] = None): self.root = Path(sequence_dir) assert self.root.exists(), f"Sequence directory not found: {self.root}" # Auto-detect format self.format = self._detect_format() print(f"[EndoDataset] Detected format: {self.format} at {self.root}") # Load file lists self.image_paths = self._find_images() if max_frames: self.image_paths = self.image_paths[:max_frames] self.depth_dir = self._find_subdir(["depth", "depths", "depth_maps"]) self.mask_dir = self._find_subdir(["masks", "mask", "tool_masks"]) self.n_frames = len(self.image_paths) # Load poses if available self.poses = self._load_poses() self.intrinsics = self._load_intrinsics() print(f"[EndoDataset] {self.n_frames} frames | " f"depth: {'✓' if self.depth_dir else '✗'} | " f"masks: {'✓' if self.mask_dir else '✗'} | " f"poses: {'✓' if self.poses is not None else '✗'}") def _detect_format(self) -> str: """Auto-detect dataset format from directory structure.""" if (self.root / "poses_bounds.npy").exists(): return "llff" elif (self.root / "color").is_dir(): return "c3vd" elif (self.root / "Frames").is_dir(): return "endoslam" elif (self.root / "images").is_dir(): return "endonerf" else: # Generic: look for image files return "generic" def _find_images(self) -> List[Path]: """Find all image files in the appropriate directory.""" search_dirs = { "llff": ["images"], "c3vd": ["color"], "endoslam": ["Frames"], "endonerf": ["images"], "generic": ["."], } for dirname in search_dirs.get(self.format, ["."]): img_dir = self.root / dirname if img_dir.is_dir(): paths = sorted( list(img_dir.glob("*.png")) + list(img_dir.glob("*.jpg")) + list(img_dir.glob("*.jpeg")) ) if paths: return paths raise FileNotFoundError(f"No images found in {self.root}") def _find_subdir(self, candidates: List[str]) -> Optional[Path]: """Find an existing subdirectory from candidates.""" for name in candidates: d = self.root / name if d.is_dir() and any(d.iterdir()): return d return None def _load_poses(self) -> Optional[np.ndarray]: """Load camera poses.""" # LLFF format pb_path = self.root / "poses_bounds.npy" if pb_path.exists(): poses_bounds = np.load(str(pb_path)) # poses_bounds: [N, 17] = [3x5 pose | near, far] poses = poses_bounds[:, :-2].reshape(-1, 3, 5) # Convert LLFF [R|t|hwf] to 4x4 c2w = np.zeros((poses.shape[0], 4, 4), dtype=np.float32) c2w[:, :3, :4] = poses[:, :3, :4] c2w[:, 3, 3] = 1.0 return c2w # JSON poses for name in ["transforms.json", "poses.json", "cameras.json"]: json_path = self.root / name if json_path.exists(): with open(json_path) as f: data = json.load(f) if "frames" in data: poses = [] for frame in data["frames"]: mat = np.array(frame["transform_matrix"], dtype=np.float32) poses.append(mat) return np.stack(poses) return None def _load_intrinsics(self) -> Optional[np.ndarray]: """Load camera intrinsics.""" # From LLFF poses pb_path = self.root / "poses_bounds.npy" if pb_path.exists(): poses_bounds = np.load(str(pb_path)) poses = poses_bounds[:, :-2].reshape(-1, 3, 5) h, w, f = poses[0, :, 4] K = np.array([[f, 0, w / 2], [0, f, h / 2], [0, 0, 1]], dtype=np.float32) return K # From transforms.json for name in ["transforms.json", "cameras.json"]: json_path = self.root / name if json_path.exists(): with open(json_path) as f: data = json.load(f) if "fl_x" in data: K = np.array([ [data["fl_x"], 0, data.get("cx", data.get("w", 640) / 2)], [0, data["fl_y"], data.get("cy", data.get("h", 480) / 2)], [0, 0, 1] ], dtype=np.float32) return K return None def __len__(self) -> int: return self.n_frames def __getitem__(self, idx: int) -> Frame: """Load a single frame with all annotations.""" from PIL import Image # RGB img = Image.open(self.image_paths[idx]).convert("RGB") rgb = np.array(img, dtype=np.uint8) # Depth depth = None if self.depth_dir: stem = self.image_paths[idx].stem for ext in [".npy", ".png", ".tiff"]: depth_path = self.depth_dir / f"{stem}{ext}" if depth_path.exists(): if ext == ".npy": depth = np.load(str(depth_path)).astype(np.float32) else: depth_img = Image.open(depth_path) depth = np.array(depth_img, dtype=np.float32) if depth.max() > 1000: depth = depth / 1000.0 # mm → m break # Mask mask = None if self.mask_dir: stem = self.image_paths[idx].stem for ext in [".png", ".jpg"]: mask_path = self.mask_dir / f"{stem}{ext}" if mask_path.exists(): mask_img = Image.open(mask_path).convert("L") mask = np.array(mask_img) > 127 # Binary break # Pose pose = None if self.poses is not None and idx < len(self.poses): pose = self.poses[idx] # Timestamp timestamp = idx / max(self.n_frames - 1, 1) return Frame( rgb=rgb, depth=depth, mask=mask, pose=pose, intrinsics=self.intrinsics, timestamp=timestamp, frame_idx=idx, sequence_name=self.root.name, ) def get_point_cloud( self, frame_indices: Optional[List[int]] = None, subsample: float = 0.001, ) -> Tuple[np.ndarray, np.ndarray]: """ Generate point cloud via depth backprojection (HGI). Implements: P = ∪_t K⁻¹ · T_t · D_t · (I_t ⊙ M_t) Args: frame_indices: Which frames to use (default: all) subsample: Fraction of points to keep (0.001 = 0.1%) Returns: (points [N, 3], colors [N, 3]) in world coordinates """ if self.intrinsics is None: raise ValueError("Cannot backproject without intrinsics") indices = frame_indices or list(range(self.n_frames)) all_points = [] all_colors = [] K_inv = np.linalg.inv(self.intrinsics) for idx in indices: frame = self[idx] if frame.depth is None: continue H, W = frame.depth.shape[:2] rgb = frame.rgb.astype(np.float32) / 255.0 # Create pixel grid u, v = np.meshgrid(np.arange(W), np.arange(H)) ones = np.ones_like(u) pixels = np.stack([u, v, ones], axis=-1).reshape(-1, 3) # [HW, 3] depth_flat = frame.depth.reshape(-1) colors_flat = rgb.reshape(-1, 3) # Valid depth mask valid = depth_flat > 1e-4 if frame.mask is not None: mask_flat = frame.mask.reshape(-1) valid = valid & mask_flat # Exclude tools pixels = pixels[valid] depth_flat = depth_flat[valid] colors_flat = colors_flat[valid] # Backproject: P_cam = D * K^-1 * [u, v, 1] cam_points = (K_inv @ pixels.T).T * depth_flat[:, None] # [N, 3] # Transform to world if pose available if frame.pose is not None: cam_points_h = np.hstack([cam_points, np.ones((len(cam_points), 1))]) world_points = (frame.pose @ cam_points_h.T).T[:, :3] else: world_points = cam_points all_points.append(world_points) all_colors.append(colors_flat) if not all_points: raise ValueError("No valid depth data found for point cloud generation") points = np.concatenate(all_points, axis=0) colors = np.concatenate(all_colors, axis=0) # Subsample if subsample < 1.0: n_keep = max(int(len(points) * subsample), 1000) indices = np.random.choice(len(points), n_keep, replace=False) points = points[indices] colors = colors[indices] print(f"[HGI] Generated point cloud: {len(points):,} points from {len(all_points)} frames") return points.astype(np.float32), colors.astype(np.float32) # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="EndoGaussian-4D Dataset Download & Organization", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--data-root", type=str, default="./data", help="Root directory for all datasets") parser.add_argument("--datasets", nargs="+", type=str, help="Dataset(s) to download (e.g., endonerf c3vd)") parser.add_argument("--list", action="store_true", help="List all available datasets") parser.add_argument("--validate", action="store_true", help="Validate existing data installation") parser.add_argument("--manifest", action="store_true", help="Generate JSON manifest of available data") args = parser.parse_args() organizer = DatasetOrganizer(args.data_root) if args.list: organizer.list_datasets() elif args.datasets: for ds in args.datasets: organizer.download(ds) organizer.validate() elif args.validate: organizer.validate() elif args.manifest: organizer.generate_manifest() else: parser.print_help() if __name__ == "__main__": main()