from __future__ import annotations import hashlib import json import threading from dataclasses import asdict, dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, TYPE_CHECKING from uuid import uuid4 from adam.dataset_lab import IMAGE_EXTENSIONS, TEXT_EXTENSIONS, VIDEO_EXTENSIONS if TYPE_CHECKING: from adam.assets import Asset, AssetRegistry DATASET_MARKERS = { "dataset_manifest.json", "metadata", "frames", "videos", "captions", "actions.jsonl", "actions.csv", } SCAN_LIMIT = 20_000 ASYNC_REFRESH_AFTER = timedelta(minutes=30) _scan_lock = threading.Lock() _active_scans: set[str] = set() def _now() -> str: return datetime.now(timezone.utc).isoformat() def _path_key(path: str | Path) -> str: resolved = str(Path(path).expanduser().resolve()) return hashlib.sha1(resolved.casefold().encode("utf-8")).hexdigest()[:16] def _atomic_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8") temporary.replace(path) def _date(value: str) -> datetime | None: try: return datetime.fromisoformat(value) except (TypeError, ValueError): return None @dataclass(slots=True) class DatasetLocation: id: str name: str path: str source: str = "user" created_at: str = field(default_factory=_now) last_seen_at: str = "" exists: bool = True @classmethod def from_dict(cls, payload: dict[str, Any]) -> "DatasetLocation": return cls( id=str(payload.get("id") or _path_key(str(payload.get("path", "")))), name=str(payload.get("name") or Path(str(payload.get("path", ""))).name or "Datasets"), path=str(payload.get("path", "")), source=str(payload.get("source") or "user"), created_at=str(payload.get("created_at") or _now()), last_seen_at=str(payload.get("last_seen_at") or ""), exists=bool(payload.get("exists", True)), ) @dataclass(slots=True) class DatasetRecord: id: str name: str path: str source: str = "asset" location_id: str = "" favorite: bool = False last_used_at: str = "" discovered_at: str = field(default_factory=_now) scanned_at: str = "" exists: bool = True item_count: int = 0 image_count: int = 0 video_count: int = 0 caption_count: int = 0 missing_caption_count: int = 0 sample_image: str = "" dataset_format: str = "Unknown" warnings: list[str] = field(default_factory=list) @classmethod def from_dict(cls, payload: dict[str, Any]) -> "DatasetRecord": return cls( id=str(payload.get("id") or _path_key(str(payload.get("path", "")))), name=str(payload.get("name") or Path(str(payload.get("path", ""))).name or "Dataset"), path=str(payload.get("path", "")), source=str(payload.get("source") or "asset"), location_id=str(payload.get("location_id") or ""), favorite=bool(payload.get("favorite", False)), last_used_at=str(payload.get("last_used_at") or ""), discovered_at=str(payload.get("discovered_at") or _now()), scanned_at=str(payload.get("scanned_at") or ""), exists=bool(payload.get("exists", True)), item_count=int(payload.get("item_count", 0) or 0), image_count=int(payload.get("image_count", 0) or 0), video_count=int(payload.get("video_count", 0) or 0), caption_count=int(payload.get("caption_count", 0) or 0), missing_caption_count=int(payload.get("missing_caption_count", 0) or 0), sample_image=str(payload.get("sample_image") or ""), dataset_format=str(payload.get("dataset_format") or "Unknown"), warnings=[str(item) for item in payload.get("warnings", []) if str(item)], ) class DatasetRegistry: """Persistent ADAM-aware index of known dataset locations and datasets.""" def __init__(self, root: Path, config: Any | None = None) -> None: self.root = root.resolve() self.path = self.root / "data" / "dataset_registry.json" self.config = config self.locations: list[DatasetLocation] = [] self.datasets: dict[str, DatasetRecord] = {} self.load() def load(self) -> None: try: payload = json.loads(self.path.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError, json.JSONDecodeError): payload = {} self.locations = [ DatasetLocation.from_dict(item) for item in payload.get("locations", []) if isinstance(item, dict) ] self.datasets = { str(key): DatasetRecord.from_dict(value) for key, value in (payload.get("datasets", {}) or {}).items() if isinstance(value, dict) } def save(self) -> None: _atomic_json( self.path, { "locations": [asdict(item) for item in self.locations], "datasets": {key: asdict(value) for key, value in self.datasets.items()}, }, ) def register_location(self, path: str | Path, *, name: str = "", source: str = "user") -> DatasetLocation: resolved = Path(path).expanduser().resolve() if not resolved.is_dir(): raise ValueError("Choose an existing dataset folder.") key = _path_key(resolved) existing = next((item for item in self.locations if item.id == key), None) if existing is None: existing = DatasetLocation( id=key, name=name.strip() or resolved.name or str(resolved), path=str(resolved), source=source, ) self.locations.insert(0, existing) else: existing.name = name.strip() or existing.name existing.path = str(resolved) existing.source = source or existing.source existing.exists = True existing.last_seen_at = _now() self.save() return existing def remove_location(self, location_id: str) -> bool: before = len(self.locations) self.locations = [item for item in self.locations if item.id != location_id] changed = len(self.locations) != before if changed: self.save() return changed def known_locations(self) -> list[DatasetLocation]: locations = list(self.locations) by_path = {Path(item.path).expanduser().resolve(): item for item in locations if item.path} for path, name, source in self._automatic_location_candidates(): try: resolved = path.expanduser().resolve() except OSError: continue if resolved in by_path: continue locations.append( DatasetLocation( id=_path_key(resolved), name=name or resolved.name or str(resolved), path=str(resolved), source=source, exists=resolved.is_dir(), last_seen_at=_now() if resolved.is_dir() else "", ) ) return locations def discover_into_assets(self, assets: "AssetRegistry", *, persist: bool = False, update_cache: bool = True) -> list["Asset"]: discovered: list[Asset] = [] records = self.discover(asset_registry=assets, refresh_missing=False, persist=update_cache) for record in records: if not record.exists: continue asset = assets.register( kind="dataset", name=record.name, path=record.path, metadata={ "dataset_registry_source": record.source, "dataset_location_id": record.location_id, }, persist=False, ) discovered.append(asset) if persist and discovered: assets.save() return discovered def discover( self, *, asset_registry: "AssetRegistry | None" = None, refresh_missing: bool = True, persist: bool = True, ) -> list[DatasetRecord]: self.load() changed = False locations = self.known_locations() known_location_ids = {location.id for location in locations} for key, record in list(self.datasets.items()): if record.source in {"adam", "tool"} and record.location_id and record.location_id not in known_location_ids: del self.datasets[key] changed = True for location in locations: exists = Path(location.path).is_dir() if location.source == "user": stored = next((item for item in self.locations if item.id == location.id), None) if stored: stored.exists = exists stored.last_seen_at = _now() if exists else stored.last_seen_at changed = True if not exists: continue for candidate in self._dataset_candidates(Path(location.path)): record = self._cached_or_sampled(candidate, source=location.source, location_id=location.id) self.datasets[record.id] = record changed = True if asset_registry is not None: for asset in getattr(asset_registry, "assets", []): if getattr(asset, "kind", "") != "dataset": continue record = self._cached_or_sampled(Path(asset.path), source="asset", location_id="") record.name = asset.name or record.name self.datasets[record.id] = record changed = True for record in self.datasets.values(): record.exists = Path(record.path).is_dir() if refresh_missing and record.exists and self._needs_refresh(record): self.refresh_async(record.path, source=record.source, location_id=record.location_id) if changed and persist: self.save() return self.sorted_records() def sorted_records(self) -> list[DatasetRecord]: records = list(self.datasets.values()) records.sort( key=lambda item: ( not item.favorite, not bool(item.last_used_at), item.last_used_at or item.discovered_at, item.name.casefold(), ), reverse=False, ) favorites = sorted([item for item in records if item.favorite], key=lambda item: item.name.casefold()) recent = sorted( [item for item in records if not item.favorite and item.last_used_at], key=lambda item: item.last_used_at, reverse=True, ) others = sorted( [item for item in records if not item.favorite and not item.last_used_at], key=lambda item: item.discovered_at, reverse=True, ) return [*favorites, *recent, *others] def record_for_path(self, path: str | Path) -> DatasetRecord: key = _path_key(path) record = self.datasets.get(key) if record is None: record = self._cached_or_sampled(Path(path), source="asset", location_id="") self.datasets[key] = record self.save() return record def favorite(self, path: str | Path, enabled: bool) -> DatasetRecord: record = self.record_for_path(path) record.favorite = bool(enabled) self.save() return record def touch(self, path: str | Path) -> DatasetRecord: record = self.record_for_path(path) record.last_used_at = _now() self.save() return record def refresh_async(self, path: str | Path, *, source: str = "asset", location_id: str = "") -> None: resolved = str(Path(path).expanduser().resolve()) key = _path_key(resolved) with _scan_lock: if key in _active_scans: return _active_scans.add(key) def worker() -> None: try: record = self._scan(Path(resolved), source=source, location_id=location_id, limit=SCAN_LIMIT) fresh = DatasetRegistry(self.root, self.config) current = fresh.datasets.get(record.id) if current: record.favorite = current.favorite record.last_used_at = current.last_used_at record.discovered_at = current.discovered_at fresh.datasets[record.id] = record fresh.save() finally: with _scan_lock: _active_scans.discard(key) threading.Thread(target=worker, name="ADAMDatasetScan", daemon=True).start() def _automatic_location_candidates(self) -> list[tuple[Path, str, str]]: candidates: list[tuple[Path, str, str]] = [(self.root / "ADAM_Datasets", "ADAM Datasets", "adam")] folders = self.config.get("tool_folders", {}) if self.config is not None else {} folders = folders if isinstance(folders, dict) else {} collector_raw = str(folders.get("dataset_collector", "")).strip() if collector_raw: candidates.append((Path(collector_raw) / "Datasets", "Dataset Collector", "tool")) for tool_id in ("ddpm_trainer", "lora_trainer", "flow_trainer", "oasis_trainer"): raw_root = str(folders.get(tool_id, "")).strip() if not raw_root: continue root = Path(raw_root) label = tool_id.removesuffix("_trainer").upper() for child in ("Datasets", "datasets", "OldDatasets"): candidates.append((root / child, f"{label} {child}", "tool")) external_tools = self.root / "config" / "external_tools.json" try: payload = json.loads(external_tools.read_text(encoding="utf-8")) except (OSError, ValueError, TypeError, json.JSONDecodeError): payload = {} for entry in payload.get("tools", []) if isinstance(payload, dict) else []: if not isinstance(entry, dict): continue backend = entry.get("backend", {}) if not isinstance(backend, dict): continue raw_root = str(backend.get("root", "")).strip() if raw_root: root = Path(raw_root) candidates.append((root / "Datasets", f"{entry.get('name', 'External tool')} Datasets", "tool")) candidates.append((root / "OldDatasets", f"{entry.get('name', 'External tool')} OldDatasets", "tool")) return candidates def _dataset_candidates(self, location: Path) -> list[Path]: candidates: list[Path] = [] try: children = [item for item in location.iterdir() if item.is_dir()] except OSError: children = [] for child in children[:1000]: if self._looks_like_dataset(child): candidates.append(child) if candidates: return candidates if self._looks_like_dataset(location): candidates.append(location) return candidates def _looks_like_dataset(self, folder: Path) -> bool: if not folder.is_dir(): return False try: names = {item.name for item in folder.iterdir()} except OSError: return False if names & DATASET_MARKERS: return True checked = 0 for item in folder.rglob("*"): if checked >= 200: break checked += 1 if item.is_file() and item.suffix.casefold() in IMAGE_EXTENSIONS | VIDEO_EXTENSIONS: return True return False def _cached_or_sampled(self, path: Path, *, source: str, location_id: str) -> DatasetRecord: key = _path_key(path) current = self.datasets.get(key) if current is not None: current.exists = path.is_dir() current.source = current.source or source current.location_id = current.location_id or location_id return current return self._scan(path, source=source, location_id=location_id, limit=800) def _needs_refresh(self, record: DatasetRecord) -> bool: if not record.scanned_at: return True scanned = _date(record.scanned_at) return scanned is None or datetime.now(timezone.utc) - scanned > ASYNC_REFRESH_AFTER def _scan(self, path: Path, *, source: str, location_id: str, limit: int) -> DatasetRecord: resolved = path.expanduser().resolve() record = DatasetRecord( id=_path_key(resolved), name=resolved.name or "Dataset", path=str(resolved), source=source, location_id=location_id, exists=resolved.is_dir(), scanned_at=_now(), ) if not resolved.is_dir(): record.warnings.append("Dataset folder is unavailable.") return record files_seen = 0 capped = False try: iterator = resolved.rglob("*") for item in iterator: if not item.is_file(): continue files_seen += 1 suffix = item.suffix.casefold() if suffix in IMAGE_EXTENSIONS: record.image_count += 1 if not record.sample_image: record.sample_image = str(item) if not any(item.with_suffix(ext).is_file() for ext in (".txt", ".caption")): record.missing_caption_count += 1 elif suffix in VIDEO_EXTENSIONS: record.video_count += 1 if suffix in TEXT_EXTENSIONS and suffix in {".txt", ".caption"}: record.caption_count += 1 if files_seen >= limit: capped = True break except OSError as exc: record.warnings.append(f"Dataset could not be scanned: {exc}") record.item_count = record.image_count + record.video_count record.dataset_format = self._format_label(resolved, record) if capped: record.warnings.append(f"Counts are sampled from the first {limit:,} files.") return record @staticmethod def _format_label(folder: Path, record: DatasetRecord) -> str: if (folder / "actions.jsonl").is_file() or (folder / "actions.csv").is_file(): return "Oasis action dataset" if record.video_count: return "Video dataset" if record.image_count and record.caption_count: return "Captioned image dataset" if record.image_count: return "Image dataset" return "Dataset folder"