app.py

#!/usr/bin/env python3
"""Standalone material/size reviewer for debris_remaining_yolo."""

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import gradio as gr
from gradio_image_annotation import image_annotator
from PIL import Image, ImageDraw


# Configuration: users only need to paste these two full paths.
IMAGES_FOLDER = Path("images")
ANNOTATIONS_FILE = Path("./dataset_04.json")

# Set to True to review only 'Not Applicable' labels
ONLY_NOT_APPLICABLE = True

# Material/size answers are saved beside the annotation JSON automatically.
ANSWERS_FILE = ANNOTATIONS_FILE.with_name(
    f"{ANNOTATIONS_FILE.stem}_material_size_answers.json"
)

MATERIAL_CHOICES = [
    "carcass", "cardboard", "clothing", "construction material",
    "displaced road equipment", "furniture/mattress", "metal",
    "plastic", "rock", "tire/rubber", "vegetation", "vehicle part",
    "wood/lumber", "unknown", "Not Applicable",
]
SIZE_CHOICES = ["small", "large"]


@dataclass(frozen=True)
class ReviewSample:
    category: str
    image_path: str
    annotation_path: str
    image_stem: str
    box_index: int
    source_class_id: int
    normalized_xywh: tuple[float, float, float, float]

    @property
    def sample_id(self) -> str:
        return f"{self.category}/{self.image_stem}__box{self.box_index:03d}"


def read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def write_json(path: Path, payload: Any) -> None:
    temporary = path.with_name(f".{path.name}.tmp")
    temporary.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    temporary.replace(path)


def parse_yolo_json_box(path: Path, image_name: str, box_index: int, instance: Any):
    if not isinstance(instance, dict) or not isinstance(instance.get("yolo"), dict):
        raise ValueError(f"{path}:{image_name}: box {box_index} has no yolo object")
    yolo = instance["yolo"]
    try:
        class_value = float(yolo["class_id"])
        class_id = int(class_value)
        box = tuple(float(yolo[field]) for field in ("x_center", "y_center", "width", "height"))
    except (KeyError, TypeError, ValueError) as exc:
        raise ValueError(f"{path}:{image_name}: invalid YOLO data for box {box_index}") from exc
    if class_value != class_id or not all(0.0 <= value <= 1.0 for value in box):
        raise ValueError(f"{path}:{image_name}: invalid YOLO values for box {box_index}")
    if box[2] <= 0 or box[3] <= 0:
        raise ValueError(f"{path}:{image_name}: box {box_index} must have positive width and height")
    return class_id, box


def collect_samples(only_not_applicable: bool = ONLY_NOT_APPLICABLE) -> list[ReviewSample]:
    if not IMAGES_FOLDER.is_dir():
        raise FileNotFoundError(f"Images folder not found: {IMAGES_FOLDER}")
    if not ANNOTATIONS_FILE.is_file():
        raise FileNotFoundError(f"Annotations file not found: {ANNOTATIONS_FILE}")

    samples = []
    category = ANNOTATIONS_FILE.stem
    payload = read_json(ANNOTATIONS_FILE)
    images = payload.get("images")
    if not isinstance(images, list):
        raise ValueError(f"Invalid annotation file: {ANNOTATIONS_FILE}")
    for image_data in images:
        if not isinstance(image_data, dict) or not isinstance(image_data.get("file_name"), str):
            raise ValueError(f"{ANNOTATIONS_FILE}: invalid image entry")
        image_path = IMAGES_FOLDER / image_data["file_name"]
        if not image_path.is_file():
            raise FileNotFoundError(f"Image listed in {ANNOTATIONS_FILE} is missing: {image_path}")
        instances = image_data.get("instances")
        if not isinstance(instances, list):
            raise ValueError(f"{ANNOTATIONS_FILE}:{image_path.name}: instances must be a list")
        source_boxes = tuple(
            parse_yolo_json_box(ANNOTATIONS_FILE, image_path.name, box_index, instance)
            for box_index, instance in enumerate(instances, start=1)
        )
        for box_index, (class_id, box) in enumerate(source_boxes, start=1):
            samples.append(ReviewSample(
                category, str(image_path.resolve()), str(ANNOTATIONS_FILE.resolve()),
                image_path.stem, box_index, class_id, box,
            ))
    if not samples:
        raise RuntimeError(f"No labeled boxes found in {ANNOTATIONS_FILE}")

    if only_not_applicable:
        answers = load_answers()
        na_samples = [
            s for s in samples
            if answers.get(s.sample_id, {}).get("material") == "Not Applicable"
        ]
        if na_samples:
            return na_samples

    return samples


def sync_from_hf_hub() -> None:
    """Download latest answers file from HF Hub on startup if running on HF."""
    import os
    import shutil
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    dataset_repo_id = os.environ.get("DATASET_REPO_ID")
    space_id = os.environ.get("SPACE_ID")

    repo_id = dataset_repo_id or space_id
    repo_type = "dataset" if dataset_repo_id else ("space" if space_id else None)

    if repo_id and repo_type:
        try:
            from huggingface_hub import hf_hub_download
            downloaded_path = hf_hub_download(
                repo_id=repo_id,
                filename=ANSWERS_FILE.name,
                repo_type=repo_type,
                token=token,
            )
            shutil.copy(downloaded_path, ANSWERS_FILE)
            print(f"Successfully downloaded latest {ANSWERS_FILE.name} from HF {repo_type}: {repo_id}")
        except Exception as exc:
            print(f"HF download attempt skipped or failed: {exc}")


_answers_synced_on_startup = False


def load_answers() -> dict[str, dict[str, Any]]:
    global _answers_synced_on_startup
    if not _answers_synced_on_startup:
        sync_from_hf_hub()
        _answers_synced_on_startup = True

    if not ANSWERS_FILE.exists():
        return {}
    payload = read_json(ANSWERS_FILE)
    samples = payload.get("samples", {})
    if not isinstance(samples, dict):
        raise ValueError(f"Invalid answers file: {ANSWERS_FILE}")
    return {str(key): value for key, value in samples.items() if isinstance(value, dict)}


def values_for(sample: ReviewSample, answers: dict[str, dict[str, Any]]):
    answer = answers.get(sample.sample_id, {})
    material = str(answer.get("material") or "") or None
    # Size is implicit when material is not applicable, so do not show an
    # unavailable option in the size dropdown.
    size = None if material == "Not Applicable" else (str(answer.get("size") or "") or None)
    return material, size


def saved_count(samples, answers) -> int:
    return sum(
        bool(answers.get(sample.sample_id, {}).get("material") and
             answers.get(sample.sample_id, {}).get("size"))
        for sample in samples
    )


def sync_to_hf_space(file_path: Path) -> str:
    """Auto-commit updated answers file back to Hugging Face Hub if running on HF."""
    import os
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    dataset_repo_id = os.environ.get("DATASET_REPO_ID")
    space_id = os.environ.get("SPACE_ID")

    if not space_id and not dataset_repo_id:
        return ""
    if not token:
        return "[Warning: HF_TOKEN missing in Space Secrets - changes saved locally only!]"

    # Preferred Method: Commit to a separate Hugging Face Dataset repository (Never restarts the Space!)
    if dataset_repo_id:
        try:
            from huggingface_hub import HfApi
            api = HfApi(token=token)
            api.upload_file(
                path_or_fileobj=str(file_path),
                path_in_repo=file_path.name,
                repo_id=dataset_repo_id,
                repo_type="dataset",
                commit_message=f"Update {file_path.name} via review app"
            )
            return f"[Committed to HF Dataset: {dataset_repo_id}]"
        except Exception as exc:
            return f"[HF Dataset Commit Failed: {exc}]"

    # Fallback Method: Commit directly to the Space repository
    if space_id:
        try:
            from huggingface_hub import HfApi
            api = HfApi(token=token)
            api.upload_file(
                path_or_fileobj=str(file_path),
                path_in_repo=file_path.name,
                repo_id=space_id,
                repo_type="space",
                commit_message=f"Update {file_path.name} via review app [skip ci]"
            )
            return f"[Committed to HF Space: {space_id}]"
        except Exception as exc:
            return f"[HF Commit Failed: {exc}]"

    return ""


def write_answers(samples, answers) -> str:
    all_samples = collect_samples(only_not_applicable=False)
    ordered = {sample.sample_id: answers[sample.sample_id] for sample in all_samples if sample.sample_id in answers}
    for key, val in answers.items():
        if key not in ordered:
            ordered[key] = val

    write_json(ANSWERS_FILE, {
        "description": "Human-reviewed material/size labels for debris_remaining_yolo",
        "images_folder": str(IMAGES_FOLDER),
        "annotations_file": str(ANNOTATIONS_FILE),
        "total_boxes": len(all_samples),
        "saved_count": sum(bool(item.get("material") and item.get("size")) for item in ordered.values()),
        "samples": ordered,
    })
    return sync_to_hf_space(ANSWERS_FILE)


def boxed_image(sample: ReviewSample):
    with Image.open(sample.image_path) as opened:
        image = opened.convert("RGB")
    width, height = image.size
    xc, yc, bw, bh = sample.normalized_xywh
    box = (
        max(0, round((xc - bw / 2) * width)),
        max(0, round((yc - bh / 2) * height)),
        min(width - 1, round((xc + bw / 2) * width)),
        min(height - 1, round((yc + bh / 2) * height)),
    )
    ImageDraw.Draw(image).rectangle(box, outline=(255, 0, 0), width=max(3, round(min(width, height) * 0.005)))
    return image


def original_image(sample: ReviewSample):
    with Image.open(sample.image_path) as opened:
        return opened.convert("RGB")


def yolo_box_to_pixels(box, width: int, height: int):
    xc, yc, bw, bh = box
    return {
        "xmin": max(0, round((xc - bw / 2) * width)),
        "ymin": max(0, round((yc - bh / 2) * height)),
        "xmax": min(width, round((xc + bw / 2) * width)),
        "ymax": min(height, round((yc + bh / 2) * height)),
    }


def image_with_all_source_boxes(sample: ReviewSample):
    """Return the original image with every source box drawn onto it."""
    image = original_image(sample)
    width, height = image.size
    draw = ImageDraw.Draw(image)
    for _, box in source_boxes_for_image(sample):
        pixels = yolo_box_to_pixels(box, width, height)
        draw.rectangle(
            (pixels["xmin"], pixels["ymin"], pixels["xmax"], pixels["ymax"]),
            outline=(255, 0, 0), width=max(3, round(min(width, height) * 0.005)),
        )
    return image


def source_boxes_for_image(sample: ReviewSample):
    """Read every current instance for this image from its category JSON file."""
    annotation_path = Path(sample.annotation_path)
    payload = read_json(annotation_path)
    for image_data in payload.get("images", []):
        if isinstance(image_data, dict) and image_data.get("file_name") == Path(sample.image_path).name:
            instances = image_data.get("instances")
            if not isinstance(instances, list):
                break
            return tuple(
                parse_yolo_json_box(annotation_path, Path(sample.image_path).name, box_index, instance)
                for box_index, instance in enumerate(instances, start=1)
            )
    raise ValueError(f"Image {Path(sample.image_path).name} is missing from {annotation_path}")


def missing_annotation(sample: ReviewSample):
    """Build the drawing value; saved boxes are baked in from the category JSON."""
    return {"image": image_with_all_source_boxes(sample), "boxes": []}


def pixels_to_yolo(box: dict[str, Any], width: int, height: int):
    xmin = max(0.0, min(float(box["xmin"]), width))
    ymin = max(0.0, min(float(box["ymin"]), height))
    xmax = max(0.0, min(float(box["xmax"]), width))
    ymax = max(0.0, min(float(box["ymax"]), height))
    if xmax <= xmin or ymax <= ymin:
        raise ValueError("Every missing-debris box must have positive width and height.")
    return (
        ((xmin + xmax) / 2) / width,
        ((ymin + ymax) / 2) / height,
        (xmax - xmin) / width,
        (ymax - ymin) / height,
    )


def append_missing_boxes_to_json(sample: ReviewSample, normalized_boxes):
    """Append new YOLO boxes to the matching image entry in its category JSON."""
    annotation_path = Path(sample.annotation_path)
    payload = read_json(annotation_path)
    target_image = next(
        (image_data for image_data in payload.get("images", [])
         if isinstance(image_data, dict) and image_data.get("file_name") == Path(sample.image_path).name),
        None,
    )
    if target_image is None or not isinstance(target_image.get("instances"), list):
        raise ValueError(f"Image is missing from {annotation_path.name}.")
    instances = target_image["instances"]
    used_numbers = {
        int(str(instance.get("box_id", ""))[3:])
        for instance in instances
        if isinstance(instance, dict) and str(instance.get("box_id", "")).startswith("box")
        and str(instance.get("box_id", ""))[3:].isdigit()
    }
    next_number = max(used_numbers, default=0) + 1
    for xc, yc, bw, bh in normalized_boxes:
        instances.append({
            "box_id": f"box{next_number:03d}",
            "yolo": {
                "class_id": sample.source_class_id,
                "x_center": round(xc, 6),
                "y_center": round(yc, 6),
                "width": round(bw, 6),
                "height": round(bh, 6),
            },
        })
        next_number += 1
    write_json(annotation_path, payload)


def validate(material: Any, size: Any):
    material, size = str(material or "").strip(), str(size or "").strip()
    if material not in MATERIAL_CHOICES:
        raise ValueError("Select a material before saving.")
    if material == "Not Applicable":
        return material, "Not Applicable"
    if size not in SIZE_CHOICES:
        raise ValueError("Select a size before saving.")
    return material, size


def build_app(samples: list[ReviewSample]):
    answers = load_answers()

    def refresh_samples():
        new_samples = collect_samples(only_not_applicable=ONLY_NOT_APPLICABLE)
        samples.clear()
        samples.extend(new_samples)

    def render(sample, mode):
        return original_image(sample) if mode == "original" else boxed_image(sample)

    def load(index, status="", mode="boxed"):
        answers.update(load_answers())
        refresh_samples()
        if not samples:
            return (None, None, None, "No remaining Not Applicable samples", "All Not Applicable samples reviewed!", 0, "Saved 0/0", mode, gr.update(value="View original"))
        index = max(0, min(int(index), len(samples) - 1))
        sample = samples[index]
        material, size = values_for(sample, answers)
        mode = "original" if mode == "original" else "boxed"
        done = saved_count(samples, answers)
        return (render(sample, mode), material, size,
                f"Instance {index + 1}/{len(samples)} | {sample.sample_id}",
                status, index, f"Saved {done}/{len(samples)}", mode,
                gr.update(value="View boxed" if mode == "original" else "View original"))

    def navigate(index, offset):
        return load(int(index) + offset)

    def jump_to(instance_number, current_index):
        """Jump using the displayed 1-based instance number."""
        try:
            number = int(instance_number)
        except (TypeError, ValueError):
            return load(current_index, "Enter an instance number from 1 to the total.")
        if not 1 <= number <= len(samples):
            return load(current_index, f"Instance must be between 1 and {len(samples)}.")
        return load(number - 1)

    def save(index, material, size, mode):
        index = max(0, min(int(index), len(samples) - 1))
        sample = samples[index]
        try:
            material, size = validate(material, size)
        except ValueError as exc:
            return load(index, f"Not saved: {exc}", mode)
        answers[sample.sample_id] = {
            "category": sample.category, "image": Path(sample.image_path).name,
            "image_path": sample.image_path, "annotation_json": sample.annotation_path,
            "box_index": sample.box_index, "source_class_id": sample.source_class_id,
            "normalized_xywh": list(sample.normalized_xywh),
            "material": material, "size": size,
        }
        hf_msg = write_answers(samples, answers)
        status_msg = f"Saved {sample.sample_id}: {material}, {size}"
        if hf_msg:
            status_msg += f" {hf_msg}"
        return load(index, status_msg, mode)

    def toggle(index, mode):
        mode = "original" if mode != "original" else "boxed"
        return render(samples[int(index)], mode), mode, gr.update(value="View boxed" if mode == "original" else "View original")

    def material_changed(material, size):
        if material == "Not Applicable":
            return None
        return size

    def toggle_missing(enabled, index):
        enabled = bool(enabled)
        value = missing_annotation(samples[int(index)]) if enabled else None
        return (
            gr.update(visible=not enabled),
            gr.update(visible=not enabled),
            gr.update(value=value, visible=enabled),
            gr.update(visible=enabled),
            gr.update(visible=enabled),
        )

    def save_missing(index, annotation):
        sample = samples[max(0, min(int(index), len(samples) - 1))]
        if not isinstance(annotation, dict):
            return missing_save_unchanged("Missing boxes were not saved: enable Missing debris first.")
        boxes = annotation.get("boxes") or []
        with Image.open(sample.image_path) as opened:
            width, height = opened.size
        try:
            normalized = [pixels_to_yolo(box, width, height) for box in boxes]
        except (KeyError, TypeError, ValueError) as exc:
            return missing_save_unchanged(f"Missing boxes were not saved: {exc}")
        if not normalized:
            return missing_save_unchanged("No new missing-debris boxes to save.")
        previous_box_count = len(source_boxes_for_image(sample))
        try:
            append_missing_boxes_to_json(sample, normalized)
        except ValueError as exc:
            return missing_save_unchanged(f"Missing boxes were not saved: {exc}")

        # Make each newly appended JSON instance available for material/size
        # review immediately, without requiring an application restart.
        samples[:] = collect_samples()
        new_box_index = previous_box_count + 1
        new_index = next(
            index for index, candidate in enumerate(samples)
            if candidate.category == sample.category
            and candidate.image_path == sample.image_path
            and candidate.box_index == new_box_index
        )
        result = list(load(
            new_index,
            f"Saved {len(normalized)} missing debris box(es) to {Path(sample.annotation_path).name}. "
            "Now label the first new box.",
            "boxed",
        ))
        result[0] = gr.update(value=boxed_image(samples[new_index]), visible=True)
        result[8] = gr.update(value="View original", visible=True)
        return (*result, False, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False))

    with gr.Blocks(
        title="Debris Material/Size Review",
        css="""
        #missing-annotator button[aria-label="Rotate counterclockwise"],
        #missing-annotator button[aria-label="Rotate clockwise"] {
            display: none !important;
        }
        """,
    ) as demo:
        index = gr.State(0)
        mode = gr.State("boxed")
        gr.Markdown("# Debris material/size review")
        gr.Markdown(
            f"Images: `{IMAGES_FOLDER}` — annotations: `{ANNOTATIONS_FILE}` — saves: `{ANSWERS_FILE}`"
        )
        image = gr.Image(label="Image — review the red box", interactive=False)
        missing_help = gr.Markdown(
            "Red boxes are existing annotations. Draw only the missing debris boxes, then save them below.",
            visible=False,
        )
        missing_annotator = image_annotator(
            label="Draw missing debris boxes",
            elem_id="missing-annotator",
            image_type="pil",
            disable_edit_boxes=True,
            sources=None,
            show_clear_button=False,
            show_download_button=False,
            show_remove_button=True,
            visible=False,
        )
        save_missing_button = gr.Button("Save missing boxes", variant="primary", visible=False)
        view_button = gr.Button("View original")
        counter = gr.Textbox(label="Current instance", interactive=False)
        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=0):
                jump_number = gr.Number(
                    label="Jump to instance #", value=1, minimum=1,
                    maximum=len(samples), precision=0,
                )
                jump_button = gr.Button("Go")
            with gr.Column(scale=1, min_width=0):
                missing_debris = gr.Checkbox(label="Missing debris", value=False)
        with gr.Row():
            material = gr.Dropdown(MATERIAL_CHOICES, label="Material")
            size = gr.Dropdown(SIZE_CHOICES, label="Size")
        status = gr.Textbox(label="Status", interactive=False)
        progress = gr.Textbox(label="Progress", interactive=False)
        with gr.Row():
            previous = gr.Button("Previous")
            save_button = gr.Button("Save", variant="primary")
            next_button = gr.Button("Next")
        outputs = [image, material, size, counter, status, index, progress, mode, view_button]
        demo.load(load, inputs=[index, status, mode], outputs=outputs)
        missing_outputs = [image, view_button, missing_annotator, missing_help, save_missing_button]
        missing_save_outputs = outputs + [missing_debris, missing_annotator, missing_help, save_missing_button]

        def missing_save_unchanged(message):
            return (
                gr.skip(), gr.skip(), gr.skip(), gr.skip(), message,
                gr.skip(), gr.skip(), gr.skip(), gr.skip(),
                gr.skip(), gr.skip(), gr.skip(), gr.skip(),
            )

        previous.click(lambda i: navigate(i, -1), inputs=index, outputs=outputs).then(
            toggle_missing, inputs=[missing_debris, index], outputs=missing_outputs)
        next_button.click(lambda i: navigate(i, 1), inputs=index, outputs=outputs).then(
            toggle_missing, inputs=[missing_debris, index], outputs=missing_outputs)
        jump_button.click(jump_to, inputs=[jump_number, index], outputs=outputs).then(
            toggle_missing, inputs=[missing_debris, index], outputs=missing_outputs)
        save_button.click(save, inputs=[index, material, size, mode], outputs=outputs).then(
            toggle_missing, inputs=[missing_debris, index], outputs=missing_outputs)
        view_button.click(toggle, inputs=[index, mode], outputs=[image, mode, view_button])
        material.change(material_changed, inputs=[material, size], outputs=size)
        missing_debris.change(
            toggle_missing, inputs=[missing_debris, index], outputs=missing_outputs)
        save_missing_button.click(
            save_missing, inputs=[index, missing_annotator], outputs=missing_save_outputs)
    return demo


def main() -> None:
    samples = collect_samples()
    print(f"Collected {len(samples)} boxes from {ANNOTATIONS_FILE}")
    print(f"Already saved: {saved_count(samples, load_answers())}/{len(samples)}")
    build_app(samples).launch(
        server_name="0.0.0.0", server_port=7860,
        allowed_paths=sorted({str(Path(sample.image_path).parent) for sample in samples}),
    )


if __name__ == "__main__":
    main()