"updated crops layout"

app.py

@@ -1,23 +1,22 @@
 """
 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.
+bounding boxes in the result image, and shows signature crops with labels in a gallery.
 """
 
 import gradio as gr
 import numpy as np
 from PIL import Image
-import matplotlib.pyplot as plt
-import io
-from typing import Tuple
+from typing import Tuple, List
 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]:
+def process_image(
+    image: np.ndarray,
+) -> Tuple[np.ndarray, str, List[Tuple[np.ndarray, str]]]:
     """
     Process an input image using the SignatureBlockModel.
 
@@ -25,7 +24,7 @@ def process_image(image: np.ndarray) -> Tuple[np.ndarray, str, np.ndarray]:
         image (np.ndarray): Input image as a numpy array.
 
     Returns:
-        Tuple[np.ndarray, str, np.ndarray]: Processed image, status, and signature crops image.
+        Tuple[np.ndarray, str, List[Tuple[np.ndarray, str]]]: Processed image, status, and list of crops with labels.
     """
     # Convert numpy array to PIL Image
     pil_image = Image.fromarray(image)
@@ -53,83 +52,50 @@ def process_image(image: np.ndarray) -> Tuple[np.ndarray, str, np.ndarray]:
     return np.array(image_with_boxes), status, signature_crops
 
 
-def resize_crop(crop: np.ndarray, factor=0.5) -> np.ndarray:
+def create_signature_crops(model: SignatureBlockModel) -> List[Tuple[np.ndarray, str]]:
     """
-    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.
+    Create individual images for each signature crop with label and score information.
 
     Args:
         model (SignatureBlockModel): The initialized SignatureBlockModel.
 
     Returns:
-        np.ndarray: Image with stacked signature crops and labels.
+        List[Tuple[np.ndarray, str]]: List of tuples containing crop images 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)
+    crop_data = []
 
-    for (ax_label, ax_image), box, label, score in zip(axes, boxes, labels, scores):
+    for box, label, score in zip(boxes, labels, scores):
         crop = model.extract_box(box)
-        crop = resize_crop(crop, 0.7)
+        # resized_crop = resize_crop(crop)
+        label_text = f"{classes[label]}, Score: {score:.2f}"
+        crop_data.append((crop, label_text))
 
-        # Set background color to black for both subplots
-        ax_label.set_facecolor("black")
-        ax_image.set_facecolor("black")
+    return crop_data
 
-        # 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")
+def resize_crop(crop: np.ndarray, max_height: int = 200) -> np.ndarray:
+    """
+    Resize a crop to a maximum height while maintaining aspect ratio.
 
-    plt.tight_layout()
+    Args:
+        crop (np.ndarray): Input crop as a numpy array.
+        max_height (int): Maximum height for the crop.
 
-    # 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)
+    Returns:
+        np.ndarray: Resized crop.
+    """
+    crop_image = Image.fromarray(crop)
+    aspect_ratio = crop_image.width / crop_image.height
+    new_height = min(crop_image.height, max_height)
+    new_width = int(new_height * aspect_ratio)
 
-    return signature_crops
+    resized_crop = crop_image.resize((new_width, new_height), Image.LANCZOS)
+    return np.array(resized_crop)
 
 
 def load_examples():
@@ -151,12 +117,22 @@ with gr.Blocks() as demo:
     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")
+        input_image = gr.Image(label="Upload Document Image", type="numpy", height=500)
+        output_image = gr.Image(label="Processed Image", type="numpy", height=500)
+
+    signature_crops = gr.Gallery(
+        label="Signature Crops",
+        show_label=True,
+        elem_id="gallery",
+        columns=[6],
+        rows=[1],
+        allow_preview=True,
+        object_fit="contain",
+        height=250,
+    )
 
     with gr.Row():
         status_box = gr.Textbox(label="Document Status")
-        signature_crops = gr.Image(label="Signature Crops")
 
     process_btn = gr.Button("Process Image")