File size: 19,280 Bytes
c61c435 f8c73f9 c61c435 f8c73f9 c61c435 f8c73f9 c61c435 f8c73f9 c61c435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 | 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"
|