Spaces:

Cnass
/

sprite

Runtime error

sprite / dataset_manager.py

Upload 11 files

2414b69 verified 2 months ago

9.31 kB

"""
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/<id>.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)

	"""
	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/<id>.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)