Create app.py

app.py

@@ -0,0 +1,458 @@
+"""
+VAMP Vision Dataset Booster — Free Playground
+
+Gradio entrypoint for Hugging Face Spaces. This application accepts one target
+object image plus an environmental context prompt, delegates dataset generation
+to the verified pipeline.py engine, packages the generated YOLO dataset into a
+single ZIP archive, and returns a visual bounding-box preview.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+import os
+import shutil
+import tempfile
+import time
+import traceback
+import zipfile
+from pathlib import Path
+from typing import Any, Dict, Iterable, Optional, Tuple
+
+import gradio as gr
+from PIL import Image, ImageDraw
+
+try:
+    import torch
+except Exception:  # pragma: no cover - torch availability depends on Space image
+    torch = None
+
+try:
+    import cv2  # noqa: F401  # Imported to ensure OpenCV is available for pipeline.py
+except Exception:  # pragma: no cover - pipeline may not require direct app-level cv2 use
+    cv2 = None
+
+try:
+    import pipeline
+except Exception as import_error:  # pragma: no cover - surfaced cleanly at runtime
+    pipeline = None
+    PIPELINE_IMPORT_ERROR = import_error
+else:
+    PIPELINE_IMPORT_ERROR = None
+
+
+APP_TITLE = "VAMP Vision Dataset Booster — Free Playground"
+APP_DESCRIPTION = (
+    "Upload 1 target object photo, input a context prompt, and download a "
+    "50-image model-ready training batch with precise YOLO bounding boxes."
+)
+
+TMP_ROOT = Path(os.getenv("VAMP_PLAYGROUND_TMP", "/tmp/vamp_playground"))
+INPUT_ROOT = Path("/tmp/inputs")
+OUTPUT_ROOT = TMP_ROOT / "outputs"
+ZIP_FILENAME = "vamp_playground_dataset.zip"
+EXPECTED_IMAGE_COUNT = 50
+PREVIEW_FILENAME = "debug_preview_0.jpg"
+
+logging.basicConfig(
+    level=os.getenv("VAMP_LOG_LEVEL", "INFO").upper(),
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+)
+logger = logging.getLogger("vamp_playground")
+
+
+class PlaygroundGenerationError(RuntimeError):
+    """Raised when the playground cannot complete dataset generation."""
+
+
+def _reset_directory(path: Path) -> None:
+    """Create an empty directory, removing stale files from previous runs."""
+    if path.exists():
+        shutil.rmtree(path, ignore_errors=True)
+    path.mkdir(parents=True, exist_ok=True)
+
+
+def _safe_prompt(context_prompt: Optional[str]) -> str:
+    """Normalize a user context prompt without introducing placeholders."""
+    prompt = (context_prompt or "").strip()
+    return prompt if prompt else "neutral studio environment with natural lighting"
+
+
+def _save_uploaded_image(input_image: Image.Image, request_dir: Path) -> Path:
+    """Validate and save the uploaded PIL image as an RGB PNG for pipeline use."""
+    if input_image is None:
+        raise PlaygroundGenerationError("Please upload a target object photo before generating a dataset.")
+
+    try:
+        image = input_image.convert("RGB")
+    except Exception as exc:
+        raise PlaygroundGenerationError(
+            "The uploaded file could not be decoded as a valid image. Please try a PNG or JPEG file."
+        ) from exc
+
+    if image.width < 8 or image.height < 8:
+        raise PlaygroundGenerationError("The uploaded image is too small. Please upload an image at least 8×8 pixels.")
+
+    INPUT_ROOT.mkdir(parents=True, exist_ok=True)
+    request_input_dir = request_dir / "inputs"
+    request_input_dir.mkdir(parents=True, exist_ok=True)
+
+    canonical_input = INPUT_ROOT / "target_object.png"
+    request_input = request_input_dir / "target_object.png"
+    image.save(canonical_input, format="PNG")
+    image.save(request_input, format="PNG")
+
+    logger.info("Saved uploaded image to %s and %s", canonical_input, request_input)
+    return request_input
+
+
+def _pipeline_function() -> Any:
+    """Return the verified pipeline execution callable or raise a helpful error."""
+    if pipeline is None:
+        raise PlaygroundGenerationError(
+            "pipeline.py could not be imported. Ensure pipeline.py is committed beside app.py in the Hugging Face Space. "
+            f"Import error: {PIPELINE_IMPORT_ERROR}"
+        )
+
+    runner = getattr(pipeline, "run_pipeline", None)
+    if runner is None or not callable(runner):
+        raise PlaygroundGenerationError("pipeline.py must expose a callable function named run_pipeline.")
+    return runner
+
+
+def _accepted_kwargs(callable_obj: Any, candidate_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Filter candidate keyword arguments to match the pipeline signature.
+
+    If pipeline.run_pipeline accepts **kwargs, pass the full production config.
+    Otherwise, only pass recognized parameters so this app remains compatible
+    with stricter verified pipeline signatures.
+    """
+    try:
+        signature = inspect.signature(callable_obj)
+    except (TypeError, ValueError):
+        return candidate_kwargs
+
+    parameters = signature.parameters
+    accepts_kwargs = any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values())
+    if accepts_kwargs:
+        return candidate_kwargs
+
+    return {key: value for key, value in candidate_kwargs.items() if key in parameters}
+
+
+def _run_core_pipeline(input_path: Path, context_prompt: str, output_dir: Path) -> Optional[Path]:
+    """Run pipeline.py with conservative CPU-safe playground parameters."""
+    runner = _pipeline_function()
+
+    cpu_device = "cpu"
+    if torch is not None:
+        try:
+            torch.set_grad_enabled(False)
+            torch.set_num_threads(max(1, min(4, os.cpu_count() or 1)))
+        except Exception:
+            logger.warning("Unable to apply torch CPU thread controls; continuing with default torch settings.")
+
+    candidate_kwargs: Dict[str, Any] = {
+        "input_image_path": str(input_path),
+        "image_path": str(input_path),
+        "target_image": str(input_path),
+        "target_object_path": str(input_path),
+        "context_prompt": context_prompt,
+        "prompt": context_prompt,
+        "environment_prompt": context_prompt,
+        "output_dir": str(output_dir),
+        "dataset_dir": str(output_dir),
+        "num_images": EXPECTED_IMAGE_COUNT,
+        "image_count": EXPECTED_IMAGE_COUNT,
+        "target_count": EXPECTED_IMAGE_COUNT,
+        "device": cpu_device,
+        "use_cpu": True,
+        "skip_lora_training": True,
+        "skip_lora": True,
+        "disable_training": True,
+        "train_lora": False,
+        "fast_mode": True,
+        "playground_mode": True,
+        "inference_steps": 8,
+        "num_inference_steps": 8,
+        "guidance_scale": 3.5,
+        "seed": 42,
+    }
+    kwargs = _accepted_kwargs(runner, candidate_kwargs)
+
+    logger.info("Starting pipeline.run_pipeline with CPU-safe playground configuration.")
+    logger.info("Pipeline output directory: %s", output_dir)
+    logger.debug("Pipeline accepted kwargs: %s", sorted(kwargs.keys()))
+
+    try:
+        result = runner(**kwargs)
+    except TypeError as first_exc:
+        logger.warning("Keyword pipeline call failed; attempting compatibility positional call: %s", first_exc)
+        try:
+            result = runner(str(input_path), context_prompt, str(output_dir))
+        except Exception as second_exc:
+            raise PlaygroundGenerationError(
+                "pipeline.run_pipeline failed with both keyword and compatibility positional calls. "
+                f"Last error: {second_exc}"
+            ) from second_exc
+    except RuntimeError as exc:
+        message = str(exc).lower()
+        if "cuda" in message or "tensor" in message or "out of memory" in message:
+            raise PlaygroundGenerationError(
+                "The generation engine reported a tensor/device configuration issue on the CPU Space. "
+                "Please verify that pipeline.py honors device='cpu' and skip_lora_training=True."
+            ) from exc
+        raise
+
+    resolved = _resolve_pipeline_output(result, output_dir)
+    logger.info("Pipeline completed. Resolved dataset directory: %s", resolved or output_dir)
+    return resolved
+
+
+def _resolve_pipeline_output(result: Any, fallback_output_dir: Path) -> Optional[Path]:
+    """Resolve the dataset directory from common pipeline return shapes."""
+    if result is None:
+        return fallback_output_dir
+
+    if isinstance(result, (str, os.PathLike)):
+        path = Path(result)
+        return path if path.exists() else fallback_output_dir
+
+    if isinstance(result, dict):
+        for key in ("dataset_dir", "output_dir", "path", "dataset_path", "root"):
+            value = result.get(key)
+            if value:
+                path = Path(value)
+                if path.exists():
+                    return path
+
+    if isinstance(result, (tuple, list)):
+        for value in result:
+            if isinstance(value, (str, os.PathLike)):
+                path = Path(value)
+                if path.exists():
+                    return path
+
+    return fallback_output_dir
+
+
+def _iter_dataset_files(dataset_dir: Path) -> Iterable[Path]:
+    """Yield generated image and YOLO label files from the dataset structure."""
+    valid_suffixes = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".txt", ".yaml", ".yml", ".json"}
+    for file_path in sorted(dataset_dir.rglob("*")):
+        if file_path.is_file() and file_path.suffix.lower() in valid_suffixes:
+            if file_path.name == ZIP_FILENAME:
+                continue
+            yield file_path
+
+
+def _find_subdir(dataset_dir: Path, name: str) -> Optional[Path]:
+    """Find a dataset subdirectory named images or labels, preferring direct children."""
+    direct = dataset_dir / name
+    if direct.is_dir():
+        return direct
+    matches = [path for path in dataset_dir.rglob(name) if path.is_dir()]
+    return matches[0] if matches else None
+
+
+def _validate_dataset_layout(dataset_dir: Path) -> None:
+    """Check that the generated dataset contains images and matching YOLO labels."""
+    images_dir = _find_subdir(dataset_dir, "images")
+    labels_dir = _find_subdir(dataset_dir, "labels")
+
+    if images_dir is None or labels_dir is None:
+        logger.warning("Expected /images and /labels folders were not both found under %s", dataset_dir)
+        return
+
+    image_files = [p for p in images_dir.rglob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".bmp"}]
+    label_files = [p for p in labels_dir.rglob("*") if p.suffix.lower() == ".txt"]
+
+    if len(image_files) < EXPECTED_IMAGE_COUNT:
+        logger.warning("Expected 50 generated images, found %d in %s", len(image_files), images_dir)
+    if len(label_files) < min(len(image_files), EXPECTED_IMAGE_COUNT):
+        logger.warning("Found %d labels for %d generated images.", len(label_files), len(image_files))
+
+    logger.info("Dataset validation heartbeat: %d images, %d labels", len(image_files), len(label_files))
+
+
+def _compile_dataset_zip(dataset_dir: Path, request_dir: Path) -> Path:
+    """Package the generated dataset into vamp_playground_dataset.zip."""
+    zip_path = request_dir / ZIP_FILENAME
+    dataset_files = list(_iter_dataset_files(dataset_dir))
+
+    if not dataset_files:
+        raise PlaygroundGenerationError(
+            "The pipeline completed but no dataset files were found. Expected generated files under /images and /labels."
+        )
+
+    logger.info("Compiling %d dataset files into %s", len(dataset_files), zip_path)
+    with zipfile.ZipFile(zip_path, mode="w", compression=zipfile.ZIP_DEFLATED, compresslevel=6) as archive:
+        for file_path in dataset_files:
+            arcname = file_path.relative_to(dataset_dir)
+            archive.write(file_path, arcname=str(arcname))
+
+    logger.info("ZIP archive ready: %s", zip_path)
+    return zip_path
+
+
+def _parse_yolo_label(label_path: Path) -> Optional[Tuple[float, float, float, float]]:
+    """Read the first YOLO box from a label file as normalized xywh values."""
+    try:
+        first_line = label_path.read_text(encoding="utf-8").strip().splitlines()[0]
+        parts = first_line.split()
+        if len(parts) < 5:
+            return None
+        _, x_center, y_center, width, height = parts[:5]
+        return float(x_center), float(y_center), float(width), float(height)
+    except Exception:
+        return None
+
+
+def _draw_preview_from_dataset(dataset_dir: Path, request_dir: Path) -> Optional[Path]:
+    """Create debug_preview_0.jpg when the pipeline did not already provide one."""
+    images_dir = _find_subdir(dataset_dir, "images")
+    labels_dir = _find_subdir(dataset_dir, "labels")
+    if images_dir is None or labels_dir is None:
+        return None
+
+    image_files = sorted(
+        p for p in images_dir.rglob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+    )
+    if not image_files:
+        return None
+
+    image_path = image_files[0]
+    label_path = labels_dir / f"{image_path.stem}.txt"
+    if not label_path.exists():
+        matching_labels = sorted(labels_dir.rglob("*.txt"))
+        label_path = matching_labels[0] if matching_labels else label_path
+
+    preview_path = request_dir / PREVIEW_FILENAME
+    with Image.open(image_path).convert("RGB") as image:
+        draw = ImageDraw.Draw(image)
+        box = _parse_yolo_label(label_path) if label_path.exists() else None
+        if box is not None:
+            x_center, y_center, box_width, box_height = box
+            img_w, img_h = image.size
+            x1 = max(0, int((x_center - box_width / 2) * img_w))
+            y1 = max(0, int((y_center - box_height / 2) * img_h))
+            x2 = min(img_w - 1, int((x_center + box_width / 2) * img_w))
+            y2 = min(img_h - 1, int((y_center + box_height / 2) * img_h))
+            line_width = max(2, img_w // 160)
+            draw.rectangle((x1, y1, x2, y2), outline=(255, 36, 36), width=line_width)
+            draw.text((x1 + 4, max(0, y1 - 18)), "YOLO box", fill=(255, 36, 36))
+        else:
+            logger.warning("No readable YOLO label found for preview image %s", image_path)
+        image.save(preview_path, format="JPEG", quality=92)
+
+    logger.info("Generated fallback preview at %s", preview_path)
+    return preview_path
+
+
+def _find_preview(dataset_dir: Path, request_dir: Path) -> Path:
+    """Return pipeline-generated debug preview or synthesize one from image/label files."""
+    candidates = [
+        dataset_dir / PREVIEW_FILENAME,
+        dataset_dir / "debug" / PREVIEW_FILENAME,
+        dataset_dir / "previews" / PREVIEW_FILENAME,
+        request_dir / PREVIEW_FILENAME,
+    ]
+    candidates.extend(dataset_dir.rglob(PREVIEW_FILENAME))
+
+    for candidate in candidates:
+        if candidate.exists() and candidate.is_file():
+            logger.info("Using bounding-box preview: %s", candidate)
+            return candidate
+
+    generated = _draw_preview_from_dataset(dataset_dir, request_dir)
+    if generated and generated.exists():
+        return generated
+
+    raise PlaygroundGenerationError(
+        "The dataset was generated, but no debug_preview_0.jpg or drawable image/label pair could be found."
+    )
+
+
+def run_playground_generation(input_image: Image.Image, context_prompt: str) -> Tuple[str, str]:
+    """
+    Gradio execution handler.
+
+    Saves the uploaded image, calls pipeline.run_pipeline with CPU-friendly
+    generation parameters, compiles generated /images and /labels into the exact
+    archive name vamp_playground_dataset.zip, and returns the preview plus ZIP.
+    """
+    started_at = time.time()
+    request_id = f"run_{int(started_at)}_{os.getpid()}"
+    request_dir = TMP_ROOT / request_id
+    output_dir = OUTPUT_ROOT / request_id
+
+    try:
+        logger.info("Generation request received: %s", request_id)
+        _reset_directory(request_dir)
+        _reset_directory(output_dir)
+
+        prompt = _safe_prompt(context_prompt)
+        input_path = _save_uploaded_image(input_image, request_dir)
+        dataset_dir = _run_core_pipeline(input_path=input_path, context_prompt=prompt, output_dir=output_dir)
+        dataset_dir = dataset_dir or output_dir
+
+        _validate_dataset_layout(dataset_dir)
+        zip_path = _compile_dataset_zip(dataset_dir=dataset_dir, request_dir=request_dir)
+        preview_path = _find_preview(dataset_dir=dataset_dir, request_dir=request_dir)
+
+        elapsed = time.time() - started_at
+        logger.info("Generation request %s completed successfully in %.2f seconds", request_id, elapsed)
+        return str(preview_path), str(zip_path)
+
+    except PlaygroundGenerationError as exc:
+        logger.error("Generation request %s failed: %s", request_id, exc)
+        raise gr.Error(str(exc)) from exc
+    except Exception as exc:
+        logger.error("Unexpected generation failure for %s: %s", request_id, exc)
+        logger.debug("Unexpected failure traceback:\n%s", traceback.format_exc())
+        raise gr.Error(
+            "Dataset generation failed due to an unexpected runtime error. "
+            "Please check the Space logs for the full traceback and verify pipeline.py output paths."
+        ) from exc
+
+
+def build_interface() -> gr.Blocks:
+    """Construct the Gradio Blocks interface for Hugging Face Spaces."""
+    with gr.Blocks(title=APP_TITLE, theme=gr.themes.Soft()) as demo:
+        gr.Markdown(f"# {APP_TITLE}")
+        gr.Markdown(APP_DESCRIPTION)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(type="pil", label="Upload Target Object Photo")
+                context_prompt = gr.Textbox(
+                    label="Environmental Context Prompt",
+                    placeholder="e.g., conveyor belt with reflections",
+                    lines=2,
+                    max_lines=4,
+                )
+                generate_button = gr.Button("Generate 50-Image Dataset", variant="primary")
+
+            with gr.Column(scale=1):
+                preview_output = gr.Image(label="Visual Smoke Test Bounding-Box Preview")
+                zip_output = gr.File(label="Download Complete 50-Image YOLO Dataset (.zip)")
+
+        generate_button.click(
+            fn=run_playground_generation,
+            inputs=[input_image, context_prompt],
+            outputs=[preview_output, zip_output],
+            api_name="generate_dataset",
+        )
+
+    return demo
+
+
+demo = build_interface()
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch()
+else:
+    demo.queue(default_concurrency_limit=1)
+