Download src/web/app.py from triumphh77/Handwritten-Text-Recognition: direct link, hf CLI and curl.
- Browser
- Download file 14.2 kB
-
https://huggingface.co/spaces/triumphh77/Handwritten-Text-Recognition/resolve/main/src/web/app.py
- Command line
-
hf download hf://spaces/triumphh77/Handwritten-Text-Recognition/src/web/app.py
-
curl -L -o app.py https://huggingface.co/spaces/triumphh77/Handwritten-Text-Recognition/resolve/main/src/web/app.py
14.2 kB
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from torchvision import transforms | |
| from PIL import Image | |
| import pandas as pd | |
| import sys | |
| import os | |
| import matplotlib.pyplot as plt | |
| # Import preprocessing and model | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))) | |
| from src.utils.preprocessing import preprocess_image, deskew | |
| from src.models.crnn import CRNN | |
| # Define device | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # Build vocabulary directly from labels.csv without loading images | |
| try: | |
| df = pd.read_csv('data/labels.csv') | |
| chars = set() | |
| for text in df['text']: | |
| if pd.notna(text): | |
| chars.update(list(str(text))) | |
| vocab = sorted(list(chars)) | |
| idx_to_char = {i+1: c for i, c in enumerate(vocab)} | |
| num_classes = len(vocab) + 1 | |
| print(f"Loaded vocabulary with {len(vocab)} characters") | |
| except Exception as e: | |
| print(f"Could not load vocabulary from labels.csv: {e}") | |
| # Fallback to standard IAM vocab if dataset not available | |
| vocab = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.,!? ") | |
| idx_to_char = {i+1: c for i, c in enumerate(vocab)} | |
| num_classes = len(vocab) + 1 | |
| # Load Model | |
| model = CRNN(img_channel=1, img_height=32, img_width=1024, num_class=num_classes).to(device) | |
| import glob | |
| def get_latest_checkpoint(weights_dir='weights'): | |
| checkpoints = glob.glob(os.path.join(weights_dir, 'crnn_baseline_epoch_*.pth')) | |
| if not checkpoints: | |
| return None | |
| # Sort by epoch number | |
| checkpoints.sort(key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0])) | |
| return checkpoints[-1] | |
| weights_path = get_latest_checkpoint() | |
| if weights_path and os.path.exists(weights_path): | |
| print(f"Loading trained weights from {weights_path}...") | |
| try: | |
| model.load_state_dict(torch.load(weights_path, map_location=device)) | |
| except Exception as e: | |
| print(f"Error loading weights perfectly (might be minor mismatch): {e}") | |
| model.load_state_dict(torch.load(weights_path, map_location=device), strict=False) | |
| else: | |
| print(f"Warning: Could not find any weights in weights/. Model will output random predictions.") | |
| model.eval() | |
| # Transform matching training exactly | |
| transform = transforms.Compose([ | |
| transforms.Resize((32, 1024)), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.5,), (0.5,)) | |
| ]) | |
| def decode_predictions(preds, idx_to_char): | |
| _, max_preds = torch.max(preds, 2) | |
| max_preds = max_preds.permute(1, 0) | |
| decoded_texts = [] | |
| for batch_idx in range(max_preds.size(0)): | |
| pred_seq = max_preds[batch_idx] | |
| decoded_seq = [] | |
| for i in range(len(pred_seq)): | |
| if pred_seq[i] != 0 and (i == 0 or pred_seq[i] != pred_seq[i-1]): | |
| char_idx = pred_seq[i].item() | |
| if char_idx in idx_to_char: | |
| decoded_seq.append(idx_to_char[char_idx]) | |
| decoded_texts.append("".join(decoded_seq)) | |
| return decoded_texts | |
| def auto_crop_image(gray_img): | |
| # Apply Gaussian blur to reduce noise | |
| blurred = cv2.GaussianBlur(gray_img, (5, 5), 0) | |
| # Apply Otsu's thresholding to separate dark ink from white background | |
| _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) | |
| # Find contours (shapes) in the image | |
| contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if not contours: | |
| return gray_img | |
| # Filter contours to exclude tiny noise and giant objects (like the pen) | |
| img_area = gray_img.shape[0] * gray_img.shape[1] | |
| valid_contours = [] | |
| for c in contours: | |
| area = cv2.contourArea(c) | |
| # Keep contours that are larger than a speck of dust but smaller than half the image | |
| if 20 < area < (img_area * 0.4): | |
| valid_contours.append(c) | |
| if not valid_contours: | |
| return gray_img # Fallback to original if filtering removes everything | |
| # Find the bounding box that encompasses all valid text contours | |
| x_min, y_min = float('inf'), float('inf') | |
| x_max, y_max = 0, 0 | |
| for c in valid_contours: | |
| x, y, w, h = cv2.boundingRect(c) | |
| x_min = min(x_min, x) | |
| y_min = min(y_min, y) | |
| x_max = max(x_max, x + w) | |
| y_max = max(y_max, y + h) | |
| # Add a generous padding around the text | |
| pad_y = int((y_max - y_min) * 0.2) | |
| pad_x = int((x_max - x_min) * 0.05) | |
| x_min = max(0, x_min - pad_x) | |
| y_min = max(0, y_min - pad_y) | |
| x_max = min(gray_img.shape[1], x_max + pad_x) | |
| y_max = min(gray_img.shape[0], y_max + pad_y) | |
| # Crop the image | |
| cropped = gray_img[y_min:y_max, x_min:x_max] | |
| # CRITICAL FIX for Out-of-Distribution aspect ratios: | |
| # The training data (IAM dataset) has an average aspect ratio of ~16:1. | |
| # The training pipeline blindly squashes images to 32x1024 (32:1 ratio). | |
| # If a user uploads a short word (like a 3:1 ratio "THANK YOU"), | |
| # it gets stretched 10x horizontally, destroying the letters! | |
| # To fix this, we pad the cropped image with white space on the right | |
| # so its aspect ratio matches the training average (16:1) BEFORE squashing. | |
| h, w = cropped.shape | |
| target_aspect_ratio = 16.0 | |
| if w / h < target_aspect_ratio: | |
| target_w = int(h * target_aspect_ratio) | |
| pad_width = target_w - w | |
| # Pad with white (255) on the right | |
| cropped = cv2.copyMakeBorder(cropped, 0, 0, 0, pad_width, cv2.BORDER_CONSTANT, value=255) | |
| return cropped | |
| def process_and_predict(image, apply_auto_crop=True): | |
| if image is None: | |
| return None, "Please upload an image.", None, None, None | |
| # Convert Gradio Image (which is a PIL Image by default) to grayscale | |
| if not isinstance(image, Image.Image): | |
| image = Image.fromarray(image) | |
| gray_image = image.convert('L') | |
| # For display purposes (Gradio output image) | |
| img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) | |
| gray_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY) | |
| # CRITICAL: Binarization (Otsu's thresholding) to force pure black text on pure white background | |
| # This removes shadows, lighting gradients, and colored paper backgrounds | |
| # that the model was never trained on. | |
| blurred = cv2.GaussianBlur(gray_cv, (5, 5), 0) | |
| _, binarized = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) | |
| if not apply_auto_crop: | |
| # If auto-crop is disabled, we bypass all fancy preprocessing to precisely | |
| # match the dataset loading behavior. This ensures dataset images work perfectly. | |
| gray_image_pil = Image.fromarray(gray_cv) | |
| img_tensor = transform(gray_image_pil).unsqueeze(0).to(device) | |
| # For display, just show what the network sees (squashed) | |
| display_processed_img = np.array(gray_image_pil.resize((1024, 32), Image.BILINEAR)) | |
| else: | |
| # Auto-crop if requested (using the binarized image for cleaner crops) | |
| processed_base = auto_crop_image(binarized) | |
| deskewed_img = deskew(processed_base) | |
| processed_img_np = preprocess_image(deskewed_img, target_size=(1024, 32)) | |
| display_processed_img = processed_img_np | |
| # Convert cropped numpy array back to PIL for tensor transform | |
| gray_image_cropped = Image.fromarray(display_processed_img) | |
| # For Model Prediction | |
| # We must use exactly the same transform as training, and pass a PIL image | |
| img_tensor = transform(gray_image_cropped).unsqueeze(0).to(device) | |
| # Predict and extract features | |
| with torch.no_grad(): | |
| # Get CNN features for activation map | |
| cnn_features = model.cnn(img_tensor) # shape: (1, 512, 1, seq_len) | |
| preds = model(img_tensor) | |
| preds = preds.permute(1, 0, 2) # (seq_len, batch, num_classes) | |
| decoded_text = decode_predictions(preds, idx_to_char)[0] | |
| # Calculate probabilities from LogSoftmax output | |
| probs = torch.exp(preds[:, 0, :]) # shape: (seq_len, num_classes) | |
| if not decoded_text.strip(): | |
| decoded_text = "[Model returned blank - Needs more training epochs]" | |
| # 1. Generate CTC Probability Matrix Heatmap | |
| probs_np = probs.cpu().numpy().T # shape: (num_classes, seq_len) | |
| fig_heatmap, ax1 = plt.subplots(figsize=(10, 4)) | |
| cax = ax1.imshow(probs_np, aspect='auto', cmap='viridis') | |
| ax1.set_title("CTC Probability Matrix Heatmap") | |
| ax1.set_xlabel("Time Frame (Sequence Steps)") | |
| ax1.set_ylabel("Vocabulary Character Index") | |
| fig_heatmap.colorbar(cax, ax=ax1, fraction=0.046, pad=0.04, label="Probability") | |
| plt.tight_layout() | |
| # 2. Generate Character Confidence Bar Chart | |
| max_probs, max_idx = torch.max(probs, dim=1) | |
| chars = [] | |
| confidences = [] | |
| for i in range(len(max_idx)): | |
| if max_idx[i] != 0 and (i == 0 or max_idx[i] != max_idx[i-1]): | |
| char_idx = max_idx[i].item() | |
| if char_idx in idx_to_char: | |
| chars.append(idx_to_char[char_idx]) | |
| confidences.append(max_probs[i].item()) | |
| # Adjust width based on number of characters | |
| fig_bar, ax2 = plt.subplots(figsize=(max(8, len(chars)*0.4), 4)) | |
| if chars: | |
| bars = ax2.bar(range(len(chars)), confidences, color='#FF9900') | |
| ax2.set_xticks(range(len(chars))) | |
| ax2.set_xticklabels(chars) | |
| ax2.set_ylim(0, 1.1) | |
| ax2.set_title("Character Confidence Scores") | |
| ax2.set_ylabel("Confidence Probability") | |
| # Add percentage labels above bars | |
| for bar in bars: | |
| yval = bar.get_height() | |
| ax2.text(bar.get_x() + bar.get_width()/2.0, yval + 0.02, | |
| f'{yval*100:.0f}%', va='bottom', ha='center', fontsize=8, rotation=45) | |
| else: | |
| ax2.text(0.5, 0.5, "No characters predicted", ha='center', va='center') | |
| plt.tight_layout() | |
| # 3. Generate CNN Feature Activation Overlay | |
| # Average the CNN features across all channels to get a 1D activation map | |
| activation = torch.mean(cnn_features, dim=1).squeeze().cpu().numpy() | |
| # Normalize activation to 0-255 | |
| activation = (activation - activation.min()) / (activation.max() - activation.min() + 1e-8) | |
| activation = (activation * 255).astype(np.uint8) | |
| # Resize to match the original image dimensions | |
| heatmap_img = cv2.resize(activation, (processed_img_np.shape[1], processed_img_np.shape[0])) | |
| # Apply color map | |
| heatmap_color = cv2.applyColorMap(heatmap_img, cv2.COLORMAP_JET) | |
| # Convert grayscale original image to BGR so we can blend it | |
| original_bgr = cv2.cvtColor(display_processed_img, cv2.COLOR_GRAY2BGR) | |
| # Overlay heatmap on original image (50% alpha blend) | |
| overlay_img = cv2.addWeighted(heatmap_color, 0.5, original_bgr, 0.5, 0) | |
| # Convert BGR to RGB for Gradio display | |
| overlay_img = cv2.cvtColor(overlay_img, cv2.COLOR_BGR2RGB) | |
| return display_processed_img, decoded_text, fig_heatmap, fig_bar, overlay_img | |
| # Redesign UI with Gradio Blocks for a proper Dashboard layout | |
| with gr.Blocks(title="Handwritten Text Recognition (HTR)", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("<h1 style='text-align: center;'>Handwritten Text Recognition (HTR) Dashboard</h1>") | |
| gr.Markdown("Upload an image of handwritten text. The system will preprocess it and extract the text using our trained custom CRNN model.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Editor tool allows manual cropping in UI before sending | |
| input_image = gr.Image(type="pil", label="Upload Handwritten Text Image") | |
| auto_crop_checkbox = gr.Checkbox(label="โจ Auto-Crop Background (Smart Vision)", value=True, info="Automatically zooms in on the text and removes giant background objects/pens.") | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear") | |
| submit_btn = gr.Button("Submit", variant="primary") | |
| with gr.Column(scale=1): | |
| output_image = gr.Image(type="numpy", label="Preprocessed (1024 x 32)") | |
| gr.Markdown("<p style='font-size: 12px; color: gray;'>Grayscale, aspect-ratio preserved, padded to 32x1024</p>") | |
| output_text = gr.Textbox(label="Predicted Text", lines=2) | |
| gr.Markdown("---") | |
| gr.Markdown("### ๐ Model Insights & Analytics (Explainable AI)") | |
| with gr.Accordion("๐ How to read these graphs (Interpretation Guide)", open=False): | |
| gr.Markdown(""" | |
| **1. CNN Feature Activation Overlay:** Shows exactly where the model's 'eyes' are focusing on the image. Red/hot areas indicate regions with strong visual features (like complex curves or sharp lines) that the Convolutional Neural Network detected. | |
| **2. CTC Probability Matrix Heatmap:** Shows *when* the model made a decision. The X-axis is the timeline (reading left-to-right), and the Y-axis contains all possible characters. Yellow dots indicate the exact moment the AI identified a specific letter. | |
| **3. Character Confidence Scores:** Shows *how sure* the model is about each letter it predicted. If the model misreads a word, this chart usually shows a low confidence score for the incorrect letter, proving it was uncertain. | |
| """) | |
| with gr.Row(): | |
| cnn_activation_image = gr.Image(type="numpy", label="1. CNN Feature Activation Overlay") | |
| with gr.Row(): | |
| heatmap_plot = gr.Plot(label="2. CTC Probability Heatmap") | |
| with gr.Row(): | |
| confidence_plot = gr.Plot(label="3. Character Confidence Scores") | |
| submit_btn.click( | |
| fn=process_and_predict, | |
| inputs=[input_image, auto_crop_checkbox], | |
| outputs=[output_image, output_text, heatmap_plot, confidence_plot, cnn_activation_image] | |
| ) | |
| clear_btn.click( | |
| fn=lambda: [None, True, None, "", None, None, None], | |
| inputs=[], | |
| outputs=[input_image, auto_crop_checkbox, output_image, output_text, heatmap_plot, confidence_plot, cnn_activation_image] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(share=True) |