Spaces:
Running
Running
File size: 5,074 Bytes
f725299 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
"""
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()
|