AareLC_train/tools/import_studio_export.py

#!/usr/bin/env python3
"""Bring a studio export (``python -m src.pipeline.run export``) into this repository.

Images once, labels per dataset: every frame lands in ``data/pool/images/`` exactly
once (LFS), and the dataset directory holds symlinks into the pool in the layout
Ultralytics expects, so labels resolve next to them:

    data/datasets/<name>/
    ├── images/{train,val}/tpi_<id>.jpg -> ../../../../pool/images/tpi_<id>.jpg
    ├── labels/{train,val}/tpi_<id>.txt
    ├── dataset.yaml          (path: . ; resolved to an absolute path at train time)
    ├── manifest.csv          (the studio's provenance: tpid, split, reason, beamline, ...)
    └── source.json           (where it came from, counts, classes, gold manifest hash)

A second dataset re-lists the frames it shares with earlier ones and adds only the
new images to the pool, so ``git diff dataset/A dataset/B`` shows what changed.

    python tools/import_studio_export.py /path/to/dataset_20261002_1200 --tag alc \\
        [--gold /path/to/studio/data/gold] [--commit] [--git-tag]
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import re
import shutil
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

ROOT = Path(__file__).resolve().parents[1]
IMG_EXT = (".jpg", ".jpeg", ".png")


def _sha256(p: Path) -> str:
    h = hashlib.sha256()
    with p.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def read_names(dataset_yaml: Path) -> Dict[int, str]:
    """``names:`` block of a YOLO dataset.yaml, without PyYAML."""
    names: Dict[int, str] = {}
    in_names = False
    for line in dataset_yaml.read_text(encoding="utf-8").splitlines():
        if re.match(r"^names\s*:", line):
            in_names = True
            continue
        if in_names:
            m = re.match(r"^\s+(\d+)\s*:\s*(\S+)", line)
            if m:
                names[int(m.group(1))] = m.group(2)
            elif line.strip() and not line.startswith(" "):
                break
    return names


def write_dataset_yaml(dest: Path, names: Dict[int, str]) -> Path:
    body = "\n".join(f"  {i}: {names[i]}" for i in sorted(names))
    p = dest / "dataset.yaml"
    p.write_text(
        "# path is resolved to this directory's absolute path by tools/train_run.py\n"
        f"path: .\ntrain: images/train\nval: images/val\n\nnames:\n{body}\n",
        encoding="utf-8",
    )
    return p


def add_to_pool(image: Path, pool: Path) -> tuple[Path, bool]:
    """Copy ``image`` into the pool unless an identical file is already there."""
    pool.mkdir(parents=True, exist_ok=True)
    dest = pool / image.name
    if dest.exists():
        if dest.stat().st_size == image.stat().st_size:
            return dest, False
        raise SystemExit(f"{dest} exists with a different size than {image}; a pool frame never changes")
    shutil.copy2(image, dest)
    return dest, True


def link_into(dataset_images: Path, pool_file: Path) -> Path:
    dataset_images.mkdir(parents=True, exist_ok=True)
    link = dataset_images / pool_file.name
    rel = os.path.relpath(pool_file, dataset_images)
    if link.is_symlink() or link.exists():
        link.unlink()
    link.symlink_to(rel)
    return link


def import_export(export_dir: Path, name: str, *, root: Path = ROOT, gold_dir: Optional[Path] = None) -> Dict:
    export_dir = Path(export_dir)
    if not (export_dir / "dataset.yaml").exists():
        raise SystemExit(f"{export_dir} is not a studio export (no dataset.yaml)")
    names = read_names(export_dir / "dataset.yaml")
    pool = root / "data" / "pool" / "images"
    dest = root / "data" / "datasets" / name
    if dest.exists():
        raise SystemExit(f"{dest} exists; datasets are immutable, pick another name")
    counts: Dict[str, Dict[str, int]] = {}
    added = 0
    for split in ("train", "val"):
        src_images = export_dir / split / "images"
        src_labels = export_dir / split / "labels"
        n_img = n_lbl = n_empty = 0
        for img in sorted(p for p in src_images.glob("*") if p.suffix.lower() in IMG_EXT):
            pool_file, new = add_to_pool(img, pool)
            added += int(new)
            link_into(dest / "images" / split, pool_file)
            lbl = src_labels / f"{img.stem}.txt"
            out_lbl = dest / "labels" / split / f"{img.stem}.txt"
            out_lbl.parent.mkdir(parents=True, exist_ok=True)
            if lbl.exists():
                shutil.copy2(lbl, out_lbl)
                n_lbl += 1
                if not lbl.read_text(encoding="utf-8").strip():
                    n_empty += 1
            else:
                out_lbl.write_text("", encoding="utf-8")  # background frame
                n_empty += 1
            n_img += 1
        counts[split] = {"images": n_img, "labels": n_lbl, "empty": n_empty}
    write_dataset_yaml(dest, names)
    if (export_dir / "manifest.csv").exists():
        shutil.copy2(export_dir / "manifest.csv", dest / "manifest.csv")
    gold_info = None
    if gold_dir:
        gold_info = import_gold(Path(gold_dir), root=root)
    source = {
        "name": name,
        "imported_at": datetime.now().isoformat(timespec="seconds"),
        "export_dir": str(export_dir.resolve()),
        "classes": names,
        "counts": counts,
        "new_pool_images": added,
        "manifest_sha256": _sha256(dest / "manifest.csv") if (dest / "manifest.csv").exists() else None,
        "gold": gold_info,
    }
    (dest / "source.json").write_text(json.dumps(source, indent=2) + "\n", encoding="utf-8")
    return source


def import_gold(gold_dir: Path, *, root: Path = ROOT) -> Dict:
    """The studio's frozen gold set: images into the pool, symlinks + labels + manifest under data/gold."""
    pool = root / "data" / "pool" / "images"
    dest = root / "data" / "gold"
    n = 0
    for img in sorted(p for p in (gold_dir / "images").glob("*") if p.suffix.lower() in IMG_EXT):
        pool_file, _ = add_to_pool(img, pool)
        link_into(dest / "images", pool_file)
        lbl = gold_dir / "labels" / f"{img.stem}.txt"
        (dest / "labels").mkdir(parents=True, exist_ok=True)
        if lbl.exists():
            shutil.copy2(lbl, dest / "labels" / lbl.name)
        n += 1
    manifest = gold_dir / "gold_manifest.json"
    info = {"images": n, "manifest_sha256": None}
    if manifest.exists():
        shutil.copy2(manifest, dest / "gold_manifest.json")
        info["manifest_sha256"] = _sha256(manifest)
    return info


def git(root: Path, *args: str) -> str:
    return subprocess.check_output(["git", "-C", str(root), *args], text=True).strip()


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("export_dir")
    ap.add_argument("--tag", required=True, help="short dataset name; directory = <date>_<tag>")
    ap.add_argument("--name", default=None, help="full dataset name instead of <date>_<tag>")
    ap.add_argument("--gold", default=None, help="studio gold directory (images/, labels/, gold_manifest.json)")
    ap.add_argument("--root", default=str(ROOT))
    ap.add_argument("--commit", action="store_true", help="git add + commit the dataset (and pool additions)")
    ap.add_argument("--git-tag", action="store_true", help="also tag the commit dataset/<name>")
    args = ap.parse_args(argv)
    root = Path(args.root)
    name = args.name or f"{datetime.now().strftime('%Y-%m-%d')}_{args.tag}"
    source = import_export(Path(args.export_dir), name, root=root, gold_dir=Path(args.gold) if args.gold else None)
    c = source["counts"]
    print(f"dataset {name}: train {c['train']['images']} / val {c['val']['images']} images, "
          f"{source['new_pool_images']} new in the pool")
    if args.commit:
        git(root, "add", "data/pool/images", f"data/datasets/{name}")
        if args.gold:
            git(root, "add", "data/gold")
        msg = (f"dataset: {name} ({c['train']['images']} train / {c['val']['images']} val, "
               f"{source['new_pool_images']} new pool images)")
        git(root, "commit", "-q", "-m", msg)
        print(f"committed: {msg}")
        if args.git_tag:
            git(root, "tag", f"dataset/{name}")
            print(f"tagged dataset/{name}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())