app.py

"""
This script creates a Gradio GUI for detecting and classifying signature blocks in document images
using the SignatureBlockModel. It loads example images from the /assets directory, displays
bounding boxes in the result image, and shows cropped signature blocks with labels in a separate view.
"""

import gradio as gr
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import io
from typing import Tuple
import os

from scripts.signature_blocks import SignatureBlockModel

ASSETS_DIR = os.path.join(os.path.dirname(__file__), "assets")


def process_image(image: np.ndarray) -> Tuple[np.ndarray, str, np.ndarray]:
    """
    Process an input image using the SignatureBlockModel.

    Args:
        image (np.ndarray): Input image as a numpy array.

    Returns:
        Tuple[np.ndarray, str, np.ndarray]: Processed image, status, and signature crops image.
    """
    # Convert numpy array to PIL Image
    pil_image = Image.fromarray(image)

    # Initialize the model
    model = SignatureBlockModel(pil_image)

    # Get processed image with boxes
    image_with_boxes = model.draw_boxes()

    # Get signature crops
    signature_crops = create_signature_crops(model)

    # Determine status
    labels = model.get_labels()
    if not labels.any():
        status = "Unsigned"
    elif all(label == 1 for label in labels):
        status = "Fully Executed"
    elif all(label == 2 for label in labels):
        status = "Unsigned"
    else:
        status = "Partially Executed"

    return np.array(image_with_boxes), status, signature_crops


def resize_crop(crop: np.ndarray, factor=0.5) -> np.ndarray:
    """
    Resize a crop to a target size.

    Args:
        crop (np.ndarray): Input crop as a numpy array.
        target_size (Tuple[int, int]): Target size for the crop.

    Returns:
        np.ndarray: Resized crop.
    """
    crop_image = Image.fromarray(crop).convert("RGB")
    crop_size = crop_image.size
    target_size = tuple(int(dim * factor) for dim in crop_size)
    print(f"Target Size: {target_size}")
    crop_image = crop_image.resize(target_size)
    return np.array(crop_image)


def create_signature_crops(model: SignatureBlockModel) -> np.ndarray:
    """
    Create an image with stacked signature crops and labels.

    Args:
        model (SignatureBlockModel): The initialized SignatureBlockModel.

    Returns:
        np.ndarray: Image with stacked signature crops and labels.
    """
    boxes = model.get_boxes()
    scores = model.get_scores()
    labels = model.get_labels()
    classes = model.classes

    # Create a figure with the correct number of subplots
    fig, axes = plt.subplots(len(boxes), 2, figsize=(10, 3 * len(boxes)))
    # plt.subplots_adjust(hspace=0.5, wspace=0.1)  # Add space between subplots

    # Ensure axes is always a 2D array, even with only one box
    if len(boxes) == 1:
        axes = axes.reshape(1, -1)

    for (ax_label, ax_image), box, label, score in zip(axes, boxes, labels, scores):
        crop = model.extract_box(box)
        crop = resize_crop(crop, 0.7)

        # Set background color to black for both subplots
        ax_label.set_facecolor("black")
        ax_image.set_facecolor("black")

        # Add label text
        label_text = f"Label: {classes[label]}\nScore: {score:.2f}"
        ax_label.text(
            0.05,
            0.5,
            label_text,
            color="white",
            fontsize=12,
            verticalalignment="center",
            horizontalalignment="left",
        )
        ax_label.axis("off")

        # Display the crop
        ax_image.imshow(crop)
        ax_image.axis("off")

    plt.tight_layout()

    # Convert the matplotlib figure to a PNG image
    buf = io.BytesIO()
    plt.savefig(buf, format="png", facecolor="black", edgecolor="none")
    buf.seek(0)
    signature_crops = np.array(Image.open(buf))
    plt.close(fig)

    return signature_crops


def load_examples():
    """
    Load example images from the /assets directory.

    Returns:
        List[List[str]]: List of example image paths.
    """
    examples = []
    for filename in os.listdir(ASSETS_DIR):
        if filename.lower().endswith((".png", ".jpg", ".jpeg")):
            examples.append([os.path.join(ASSETS_DIR, filename)])
    return examples


with gr.Blocks() as demo:
    gr.Markdown("# Signature Block Detection")
    gr.Markdown("Upload a document image to detect and classify signature blocks.")

    with gr.Row():
        input_image = gr.Image(label="Upload Document Image")
        output_image = gr.Image(label="Processed Image")

    with gr.Row():
        status_box = gr.Textbox(label="Document Status")
        signature_crops = gr.Image(label="Signature Crops")

    process_btn = gr.Button("Process Image")

    examples = gr.Examples(
        examples=load_examples(),
        inputs=input_image,
    )

    process_btn.click(
        fn=process_image,
        inputs=input_image,
        outputs=[output_image, status_box, signature_crops],
    )

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