Spaces:
Paused
Paused
| """Load local image folders or video files into UnifiedSequence.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import numpy as np | |
| from sensorflow.adapters.base import VendorAdapter | |
| from sensorflow.schemas.unified_frame import ( | |
| CameraView, | |
| EgoPose, | |
| FusedFrame, | |
| LidarData, | |
| UnifiedSequence, | |
| ) | |
| IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} | |
| VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v"} | |
| # High enough for full demo sequences; callers can raise or set None for uncapped. | |
| DEFAULT_MAX_FRAMES = 10_000 | |
| def discover_media( | |
| path: Path, | |
| *, | |
| dataset_type: Optional[str] = None, | |
| ) -> Tuple[List[Path], List[Path]]: | |
| """Return (image_files, video_files) under path (file or directory).""" | |
| from sensorflow.av_media_filter import ( | |
| IMAGE_EXTS as AV_IMAGE_EXTS, | |
| VIDEO_EXTS as AV_VIDEO_EXTS, | |
| partition_media_files, | |
| ) | |
| if path.is_file(): | |
| ext = path.suffix.lower() | |
| if ext in AV_IMAGE_EXTS: | |
| images, _ = partition_media_files([path], root=path.parent, dataset_type=dataset_type) | |
| return images, [] | |
| if ext in AV_VIDEO_EXTS: | |
| _, excluded = partition_media_files([path], root=path.parent, dataset_type=dataset_type, allow_videos=True) | |
| return [], [] if excluded else [path] | |
| return [], [] | |
| if not path.is_dir(): | |
| return [], [] | |
| raw_images: List[Path] = [] | |
| raw_videos: List[Path] = [] | |
| for p in sorted(path.rglob("*")): | |
| if not p.is_file(): | |
| continue | |
| ext = p.suffix.lower() | |
| if ext in AV_IMAGE_EXTS: | |
| raw_images.append(p) | |
| elif ext in AV_VIDEO_EXTS: | |
| raw_videos.append(p) | |
| images, _ = partition_media_files(raw_images, root=path, dataset_type=dataset_type) | |
| videos, _ = partition_media_files( | |
| raw_videos, | |
| root=path, | |
| dataset_type=dataset_type, | |
| allow_videos=True, | |
| ) | |
| return images, videos | |
| def _write_placeholder_lidar(path: Path, seed: int = 0) -> int: | |
| """Minimal LiDAR cloud so downstream 3D lift still has points.""" | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| rng = np.random.default_rng(seed) | |
| points = rng.normal(0.0, 2.0, size=(128, 3)).astype(np.float32) | |
| points.tofile(path) | |
| return len(points) | |
| def _decode_video_frames( | |
| video_path: Path, | |
| output_dir: Path, | |
| max_frames: Optional[int], | |
| start_index: int = 0, | |
| ) -> List[Path]: | |
| """Decode video to JPEG frames under output_dir. Returns written paths.""" | |
| try: | |
| import cv2 | |
| except ImportError as exc: | |
| raise RuntimeError( | |
| "opencv-python is required to decode video files for local ingest" | |
| ) from exc | |
| cap = cv2.VideoCapture(str(video_path)) | |
| if not cap.isOpened(): | |
| raise RuntimeError(f"Could not open video: {video_path}") | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| written: List[Path] = [] | |
| idx = 0 | |
| limit = max_frames if max_frames is not None else DEFAULT_MAX_FRAMES | |
| try: | |
| while len(written) < limit: | |
| ok, frame = cap.read() | |
| if not ok: | |
| break | |
| out = output_dir / f"vid_{start_index + idx:06d}.jpg" | |
| cv2.imwrite(str(out), frame) | |
| written.append(out) | |
| idx += 1 | |
| finally: | |
| cap.release() | |
| return written | |
| class LocalSequenceAdapter(VendorAdapter): | |
| """Normalize a local image directory or video into UnifiedSequence.""" | |
| def load(self, source: Dict[str, Any], sequence_id: str) -> UnifiedSequence: | |
| raw_path = source.get("source_path") or source.get("path") or "data" | |
| path = Path(str(raw_path)).expanduser() | |
| max_frames = source.get("max_frames", DEFAULT_MAX_FRAMES) | |
| if max_frames is not None: | |
| max_frames = int(max_frames) | |
| if not path.exists(): | |
| raise FileNotFoundError( | |
| f"Local source path not found: {path}. " | |
| "Set Dataset Configuration → Images Path to a folder of frames " | |
| "or a video file, then re-run ingest with Local enabled." | |
| ) | |
| images, videos = discover_media(path, dataset_type=source.get("dataset_type")) | |
| frame_images = list(images) | |
| decode_dir = Path("runs/pipeline") / sequence_id / "decoded_frames" | |
| for i, video in enumerate(videos): | |
| remaining = None if max_frames is None else max(0, max_frames - len(frame_images)) | |
| if remaining == 0: | |
| break | |
| decoded = _decode_video_frames( | |
| video, | |
| decode_dir / video.stem, | |
| remaining, | |
| start_index=len(frame_images), | |
| ) | |
| frame_images.extend(decoded) | |
| if not frame_images: | |
| raise FileNotFoundError( | |
| f"No images or videos found under {path}. " | |
| f"Supported images: {sorted(IMAGE_EXTS)}; videos: {sorted(VIDEO_EXTS)}." | |
| ) | |
| if max_frames is not None: | |
| frame_images = frame_images[:max_frames] | |
| lidar_dir = Path("runs/pipeline") / sequence_id / "lidar" | |
| frames: List[FusedFrame] = [] | |
| for i, img_path in enumerate(frame_images): | |
| frame_id = f"local_{i:06d}" | |
| lidar_path = lidar_dir / f"{frame_id}_lidar.bin" | |
| num_points = _write_placeholder_lidar(lidar_path, seed=i) | |
| frames.append( | |
| FusedFrame( | |
| frame_id=frame_id, | |
| timestamp_us=i * 100_000, | |
| lidar=LidarData(path=str(lidar_path), num_points=num_points), | |
| cameras={"front": CameraView(image_path=str(img_path.resolve()))}, | |
| ego_pose=EgoPose(speed_kmh=0.0, x=float(i) * 0.5), | |
| ground_truth=[], | |
| ) | |
| ) | |
| return UnifiedSequence( | |
| sequence_id=sequence_id, | |
| vendor="local", | |
| frames=frames, | |
| taxonomy_manifest={ | |
| "source_path": str(path.resolve()), | |
| "demo_stub": False, | |
| "total_frames": len(frames), | |
| "images_discovered": len(images), | |
| "videos_discovered": len(videos), | |
| "max_frames": max_frames, | |
| }, | |
| ) | |