app.py

import tempfile
import zipfile
from pathlib import Path
from uuid import uuid4

import cv2
import gradio as gr
import numpy as np

from cleaning import (
    clean_display_image_and_mask,
    cleaned_image_name,
    read_image_preview_rgb,
    read_image_rgb_and_preview,
    write_image_rgb,
)

_BATCH_OUTPUTS = tempfile.TemporaryDirectory(prefix="ihc-cleaner-")
_EXAMPLE_IMAGES = [["data/example.png"], ["data/example.tiff"]]
_IMAGE_FILE_TYPES = ["image", ".tif", ".tiff"]
_PREVIEW_MAX_DIMENSION = 1200
_MASK_OUTLINE_KERNEL_SIZE = 3


def _uploaded_path(file_path: str | Path | None) -> Path:
    if not file_path:
        raise ValueError("Upload an input image.")

    return Path(file_path)


def _preview_image(image):
    height, width = image.shape[:2]
    largest_dimension = max(height, width)

    if largest_dimension <= _PREVIEW_MAX_DIMENSION:
        return image

    scale = _PREVIEW_MAX_DIMENSION / largest_dimension
    preview_size = (round(width * scale), round(height * scale))
    return cv2.resize(image, preview_size, interpolation=cv2.INTER_AREA)


def _mask_outline(image: np.ndarray, debris_mask: np.ndarray) -> np.ndarray:
    outline_kernel = np.ones(
        (_MASK_OUTLINE_KERNEL_SIZE, _MASK_OUTLINE_KERNEL_SIZE),
        dtype=np.uint8,
    )
    boundary = cv2.morphologyEx(
        debris_mask,
        cv2.MORPH_GRADIENT,
        outline_kernel,
    ) > 0
    outlined = image.copy()
    outlined[boundary] = [0, 255, 0]
    return outlined


def preview_image(file_path: str | Path | None):
    if not file_path:
        return None, None, None, None, None

    return _preview_image(read_image_preview_rgb(file_path)), None, None, None, None


def clean_uploaded_image(file_path: str | Path | None):
    image_path = _uploaded_path(file_path)
    output_dir = Path(_BATCH_OUTPUTS.name) / uuid4().hex
    output_dir.mkdir(parents=True)

    image_rgb, display_rgb = read_image_rgb_and_preview(image_path)
    cleaned_display_rgb, debris_mask = clean_display_image_and_mask(
        image_rgb,
        display_rgb,
    )
    input_preview = _preview_image(display_rgb)
    preview_mask = debris_mask
    if preview_mask.shape != input_preview.shape[:2]:
        preview_mask = cv2.resize(
            preview_mask,
            (input_preview.shape[1], input_preview.shape[0]),
            interpolation=cv2.INTER_NEAREST,
        )
    mask_outline = _mask_outline(input_preview, preview_mask)
    cleaned_path = output_dir / cleaned_image_name(image_path)
    mask_path = output_dir / f"{image_path.stem}_debris_mask.png"
    write_image_rgb(cleaned_path, cleaned_display_rgb)
    if not cv2.imwrite(str(mask_path), debris_mask):
        raise OSError(f"Could not write debris mask: {mask_path}")

    return (
        input_preview,
        mask_outline,
        _preview_image(cleaned_display_rgb),
        str(mask_path),
        str(cleaned_path),
    )


def clean_directory(image_paths: list[str] | None) -> tuple[str, str]:
    if not image_paths:
        raise ValueError("Upload a directory containing images.")

    batch_dir = Path(_BATCH_OUTPUTS.name) / uuid4().hex
    cleaned_dir = batch_dir / "cleaned"
    cleaned_dir.mkdir(parents=True)

    used_names: set[str] = set()

    for image_path_string in image_paths:
        image_path = Path(image_path_string)
        image_rgb, display_rgb = read_image_rgb_and_preview(image_path)
        cleaned_rgb, _ = clean_display_image_and_mask(image_rgb, display_rgb)

        output_name = cleaned_image_name(image_path, used_names)
        write_image_rgb(cleaned_dir / output_name, cleaned_rgb)

    archive_path = batch_dir / "cleaned_images.zip"
    with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as archive:
        for output_path in sorted(cleaned_dir.iterdir()):
            archive.write(output_path, arcname=output_path.name)

    count = len(used_names)
    return (
        str(archive_path),
        f"Created cleaned images for {count} "
        f"image{'s' if count != 1 else ''}.",
    )


def build_app() -> gr.Blocks:
    with gr.Blocks(title="Image Cleanup") as app:
        gr.Markdown(
            "# Image Cleanup\n"
            "Upload an input image to remove dark, low-saturation artifacts."
        )

        with gr.Tab("Single image"):
            with gr.Row():
                input_image = gr.File(
                    label="Input image",
                    file_types=_IMAGE_FILE_TYPES,
                    type="filepath",
                )
                mask_file = gr.File(
                    label="Debris mask",
                    interactive=False,
                )
                cleaned_file = gr.File(
                    label="Cleaned image file",
                    interactive=False,
                )
            with gr.Row():
                input_preview = gr.Image(
                    label="Input preview",
                    format="png",
                    buttons=["fullscreen"],
                    interactive=False,
                )
                outline_preview = gr.Image(
                    label="Mask outline (green = boundary)",
                    format="png",
                    buttons=["fullscreen"],
                    interactive=False,
                )
                cleaned_preview = gr.Image(
                    label="Cleaned image preview",
                    format="png",
                    buttons=["fullscreen"],
                    interactive=False,
                )

            with gr.Row():
                clean_button = gr.Button(
                    "Clean image",
                    variant="primary",
                )
                gr.ClearButton(
                    [
                        input_image,
                        input_preview,
                        outline_preview,
                        cleaned_preview,
                        mask_file,
                        cleaned_file,
                    ]
                )

            gr.Examples(
                examples=_EXAMPLE_IMAGES,
                inputs=input_image,
                label="Example image",
            )

            input_image.change(
                fn=preview_image,
                inputs=input_image,
                outputs=[
                    input_preview,
                    outline_preview,
                    cleaned_preview,
                    mask_file,
                    cleaned_file,
                ],
                api_name=False,
            )

            clean_button.click(
                fn=clean_uploaded_image,
                inputs=input_image,
                outputs=[
                    input_preview,
                    outline_preview,
                    cleaned_preview,
                    mask_file,
                    cleaned_file,
                ],
                api_name="clean_image",
            )

        with gr.Tab("Image directory"):
            directory_input = gr.File(
                label="Image directory",
                file_count="directory",
                file_types=_IMAGE_FILE_TYPES,
                type="filepath",
            )
            clean_directory_button = gr.Button(
                "Clean directory",
                variant="primary",
            )
            batch_status = gr.Textbox(label="Status", interactive=False)
            batch_download = gr.File(
                label="Cleaned images",
                interactive=False,
            )

            clean_directory_button.click(
                fn=clean_directory,
                inputs=directory_input,
                outputs=[batch_download, batch_status],
                api_name="clean_directory",
            )

    return app


demo = build_app()


if __name__ == "__main__":
    demo.launch(ssr_mode=False)