""" dataset_manager.py -------------------- Staging system + batch push to the Hugging Face Dataset repo. Critical rule from the spec: do NOT commit to the HF repo once per accepted sprite. Sprites accumulate in a local staging area (`data/staging` for pending, `data/accepted` for user-accepted-but-not-yet-pushed), and only the explicit "UPLOAD TO DATASET" action performs ONE batch commit that includes every currently-accepted, not-yet-uploaded sprite. The huggingface_hub dependency is imported lazily inside the upload function so the rest of the app (segmentation, review UI, etc.) works completely fine even before `huggingface_hub` is installed / before HF_TOKEN is configured -- only the actual upload step needs it. """ from __future__ import annotations import hashlib import json import shutil import time from dataclasses import dataclass from pathlib import Path import database as db from config import ( ACCEPTED_DIR, HF_DATASET_REPO, HF_TOKEN, REJECTED_DIR, SPLIT_BY_SHEET, SPLIT_SEED, SPLIT_TRAIN_RATIO, SPLIT_VALIDATION_RATIO, STAGING_DIR, ) @dataclass class UploadResult: success: bool uploaded_count: int commit_url: str | None error: str | None split_counts: dict | None = None # --------------------------------------------------------------------------- # Train / validation / test split # --------------------------------------------------------------------------- def assign_split(sprite_id: str, sheet_id: str) -> str: """ Deterministically place one sprite in train / validation / test. Two properties matter here and neither is optional: - **Deterministic.** The split comes from a hash of the grouping key, not from a shuffle, so re-running an upload, adding sprites later, or rebuilding the dataset from scratch all produce the same assignment. A split that moves between runs silently invalidates every evaluation number you have. - **Grouped by sheet** (see SPLIT_BY_SHEET). An animation's frames are near-identical images; putting frame 7 in train and frame 8 in test measures memorisation, not generalisation. Grouping the whole sheet keeps that leak out. """ key = sheet_id if SPLIT_BY_SHEET else sprite_id digest = hashlib.sha256(f"{SPLIT_SEED}:{key}".encode("utf-8")).hexdigest() # Top 8 hex digits as a uniform 0..1 value. position = int(digest[:8], 16) / 0xFFFFFFFF if position < SPLIT_TRAIN_RATIO: return "train" if position < SPLIT_TRAIN_RATIO + SPLIT_VALIDATION_RATIO: return "validation" return "test" def staging_path_for(sprite_id: str, transparent: bool = False) -> Path: suffix = "_transparent" if transparent else "" return STAGING_DIR / f"{sprite_id}{suffix}.png" def accepted_path_for(sprite_id: str, transparent: bool = False) -> Path: suffix = "_transparent" if transparent else "" return ACCEPTED_DIR / f"{sprite_id}{suffix}.png" def rejected_path_for(sprite_id: str, transparent: bool = False) -> Path: suffix = "_transparent" if transparent else "" return REJECTED_DIR / f"{sprite_id}{suffix}.png" def move_to_accepted(sprite_id: str) -> None: """Move a sprite's staged PNG files into the accepted/ directory.""" for transparent in (False, True): src = staging_path_for(sprite_id, transparent) if src.exists(): dst = accepted_path_for(sprite_id, transparent) shutil.move(str(src), str(dst)) db.update_sprite_status(sprite_id, "accepted") def move_to_rejected(sprite_id: str) -> None: """Move a sprite's staged PNG files into rejected/ (kept for audit, not deleted).""" for transparent in (False, True): src = staging_path_for(sprite_id, transparent) if src.exists(): dst = rejected_path_for(sprite_id, transparent) shutil.move(str(src), str(dst)) db.update_sprite_status(sprite_id, "rejected") def _row_value(row, key, default=None): """Read a column that may be absent on rows written by an older schema.""" try: value = row[key] except (IndexError, KeyError): return default return default if value is None else value def build_manifest(sprite_rows: list, splits: dict[str, str] | None = None) -> list[dict]: """Build the JSONL-ready metadata list for every row in `sprite_rows`.""" splits = splits or {} manifest = [] for row in sprite_rows: quality = json.loads(row["quality_json"]) if row["quality_json"] else {} reasons = json.loads(row["confidence_reasons"]) if row["confidence_reasons"] else [] tags = json.loads(row["tags"]) if row["tags"] else [] manifest.append( { "id": row["id"], "file_name": f"images/{row['id']}.png", "source_file": row["source_file"], "source_bbox": { "x": row["x"], "y": row["y"], "width": row["width"], "height": row["height"], }, "width": row["width"], "height": row["height"], "category": row["category"], "style": row["style"], "tags": tags, "quality_score": quality.get("overall_quality_score"), "pixel_art_probability": quality.get("pixel_art_probability"), "has_alpha": quality.get("has_alpha"), "unique_colors": quality.get("unique_colors"), "sha256": row["sha256"], "phash": row["phash"], "segmentation_confidence": row["confidence"], "segmentation_reasons": reasons, "segmentation_source": _row_value(row, "segmentation_source", "gap"), # True while the category/style/tags on this row are still # a machine suggestion nobody reviewed -- consumers that # need human-verified labels should filter on it. "auto_tagged": bool(_row_value(row, "auto_tagged", 0)), "split": splits.get(row["id"], _row_value(row, "split", "train")), "approved": True, } ) return manifest def push_accepted_batch(repo_id: str | None = None, token: str | None = None) -> UploadResult: """ Push every currently-'accepted' (not yet 'uploaded') sprite to the HF Dataset repo in a SINGLE commit: images/.png for every sprite plus one metadata.jsonl describing the whole batch. This is the only place in the whole codebase that talks to the Hugging Face Hub. """ repo_id = repo_id or HF_DATASET_REPO token = token or HF_TOKEN if not token: return UploadResult(False, 0, None, "HF_TOKEN is not set. Add it as a Space secret before uploading.") if not repo_id: return UploadResult(False, 0, None, "HF_DATASET_REPO is not set.") rows = db.list_sprites(status="accepted") if not rows: return UploadResult(True, 0, None, None) try: from huggingface_hub import CommitOperationAdd, HfApi except ImportError: return UploadResult(False, 0, None, "huggingface_hub is not installed. Add it to requirements.txt.") # Assign (and persist) a split for anything that doesn't have one yet, # so the same sprite always lands in the same split no matter how many # times the dataset is rebuilt. splits: dict[str, str] = {} for row in rows: existing = _row_value(row, "split", None) split = existing or assign_split(row["id"], row["sheet_id"]) splits[row["id"]] = split if not existing: db.update_sprite_fields(row["id"], split=split) manifest = build_manifest(rows, splits=splits) operations = [] for row, entry in zip(rows, manifest): png_path = accepted_path_for(row["id"]) if not png_path.exists(): # Sprite record exists but the file went missing -- skip it # rather than failing the whole batch, and note it. continue operations.append( CommitOperationAdd(path_in_repo=entry["file_name"], path_or_fileobj=str(png_path)) ) jsonl_lines = "\n".join(json.dumps(m, ensure_ascii=False) for m in manifest) batch_filename = f"metadata/batch_{int(time.time())}.jsonl" operations.append( CommitOperationAdd( path_in_repo=batch_filename, path_or_fileobj=jsonl_lines.encode("utf-8"), ) ) api = HfApi(token=token) try: commit_info = api.create_commit( repo_id=repo_id, repo_type="dataset", operations=operations, commit_message=f"Add {len(rows)} sprites (batch upload, {time.strftime('%Y-%m-%d %H:%M:%S')})", ) except Exception as exc: # noqa: BLE001 - surface any Hub error to the UI return UploadResult(False, 0, None, str(exc)) for row in rows: db.update_sprite_status(row["id"], "uploaded") split_counts: dict[str, int] = {} for split in splits.values(): split_counts[split] = split_counts.get(split, 0) + 1 commit_url = getattr(commit_info, "commit_url", None) return UploadResult(True, len(rows), commit_url, None, split_counts=split_counts)