90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib, json
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
MODALITIES = {
|
|
"featurescript": ("featurescript_rp", ".txt"),
|
|
"step": ("step_abc", ".step"),
|
|
"stl": ("stl_abc", ".stl"),
|
|
"image": ("multiview_images_abc", ".png"),
|
|
"annotation": ("text_annotations", ".txt"),
|
|
}
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
@dataclass
|
|
class Sample:
|
|
sample_id: str
|
|
files: dict[str, str] = field(default_factory=dict)
|
|
hashes: dict[str, str] = field(default_factory=dict)
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
diagnostics: list[dict[str, Any]] = field(default_factory=list)
|
|
|
|
def as_dict(self) -> dict[str, Any]: return asdict(self)
|
|
|
|
|
|
def _files(root: Path, directory: str, suffix: str) -> dict[str, Path]:
|
|
base = root / directory; result: dict[str, Path] = {}
|
|
if not base.exists(): return result
|
|
for path in base.rglob(f"*{suffix}"):
|
|
if path.is_file() and path.stem.isdigit(): result[path.stem] = path
|
|
return result
|
|
|
|
|
|
def _jsonl_records(path: Path) -> list[dict[str, Any]]:
|
|
if not path.exists(): return []
|
|
records = []
|
|
for index, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
|
|
try:
|
|
value = json.loads(line); value["_line"] = index; records.append(value)
|
|
except json.JSONDecodeError:
|
|
records.append({"_line": index, "_invalid": True})
|
|
return records
|
|
|
|
|
|
def _assistant_script(record: dict[str, Any]) -> str | None:
|
|
for message in record.get("messages") or []:
|
|
if message.get("role") == "assistant" and isinstance(message.get("content"), str) and "FeatureScript" in message["content"]:
|
|
return message["content"]
|
|
return None
|
|
|
|
|
|
def scan_dataset(root: Path, *, include_hashes: bool = True) -> list[Sample]:
|
|
if not root.is_dir(): raise FileNotFoundError(f"CADFS input directory does not exist: {root}")
|
|
by_modality = {name: _files(root, directory, suffix) for name, (directory, suffix) in MODALITIES.items()}
|
|
ids = sorted(set().union(*(set(values) for values in by_modality.values())))
|
|
text_records = _jsonl_records(root / "CADFS_text_test.jsonl")
|
|
image_records = _jsonl_records(root / "CADFS_image_test.jsonl")
|
|
text_by_hash = {hashlib.sha256(script.encode()).hexdigest(): record for record in text_records if (script := _assistant_script(record))}
|
|
samples: list[Sample] = []
|
|
for index, sample_id in enumerate(ids):
|
|
sample = Sample(sample_id)
|
|
for name, values in by_modality.items():
|
|
path = values.get(sample_id)
|
|
if path is None: sample.diagnostics.append({"code": "missing_modality", "modality": name}); continue
|
|
sample.files[name] = str(path)
|
|
if include_hashes: sample.hashes[name] = sha256(path)
|
|
fs_path = by_modality["featurescript"].get(sample_id)
|
|
if fs_path:
|
|
script_hash = hashlib.sha256(fs_path.read_text(encoding="utf-8").encode()).hexdigest()
|
|
record = text_by_hash.get(script_hash)
|
|
if record:
|
|
sample.metadata["text_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
|
|
elif index < len(text_records):
|
|
fallback = text_records[index]; sample.metadata["text_jsonl"] = {"line": fallback.get("_line"), "cad_file_id": fallback.get("cad_file_id"), "alignment": "fallback"}
|
|
sample.diagnostics.append({"code": "alignment_fallback", "modality": "text_jsonl"})
|
|
if index < len(image_records):
|
|
record = image_records[index]
|
|
sample.metadata["image_jsonl"] = {"line": record.get("_line"), "cad_file_id": record.get("cad_file_id")}
|
|
samples.append(sample)
|
|
return samples
|