#!/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()