Spaces:
Sleeping
Sleeping
File size: 2,163 Bytes
182efca | 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 | """Sample dataset discovery for the public example loader."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
_SUFFIX_EXTERIOR = "_exterior"
_SUFFIX_INTERIOR = "_interior"
_ALLOWED_EXT = (".step", ".stp")
@dataclass(frozen=True)
class SampleFile:
filename: str
path: Path
size: int
@dataclass(frozen=True)
class SamplePair:
id: str
name: str
exterior: SampleFile
interior: SampleFile
def _pretty_name(stem: str) -> str:
return stem.replace("_", " ").replace("-", " ").strip().title()
def _stem_for(path: Path, suffix: str) -> str | None:
stem = path.stem
return stem[: -len(suffix)] if stem.endswith(suffix) else None
def discover_samples(samples_dir: Path) -> list[SamplePair]:
"""Scan `samples_dir` for `{id}_exterior.(step|stp)` + `{id}_interior.(step|stp)` pairs."""
if not samples_dir.exists() or not samples_dir.is_dir():
return []
by_id_ext: dict[str, Path] = {}
by_id_int: dict[str, Path] = {}
for file in sorted(samples_dir.iterdir()):
if not file.is_file() or file.suffix.lower() not in _ALLOWED_EXT:
continue
if (sid := _stem_for(file, _SUFFIX_EXTERIOR)) is not None:
by_id_ext[sid] = file
elif (sid := _stem_for(file, _SUFFIX_INTERIOR)) is not None:
by_id_int[sid] = file
pairs: list[SamplePair] = []
for sid in sorted(set(by_id_ext) & set(by_id_int)):
ext_path = by_id_ext[sid]
int_path = by_id_int[sid]
pairs.append(
SamplePair(
id=sid,
name=_pretty_name(sid),
exterior=SampleFile(
filename=ext_path.name, path=ext_path, size=ext_path.stat().st_size
),
interior=SampleFile(
filename=int_path.name, path=int_path, size=int_path.stat().st_size
),
)
)
return pairs
def get_sample(samples_dir: Path, sample_id: str) -> SamplePair | None:
for pair in discover_samples(samples_dir):
if pair.id == sample_id:
return pair
return None
|