Create App.py

app.py

@@ -0,0 +1,342 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from PIL import Image, ImageFont, ImageDraw
+import numpy as np
+import os
+import string
+import cv2
+from torchvision.transforms.functional import to_pil_image
+import matplotlib.pyplot as plt
+import math
+
+# --------- Globals --------- #
+CHARS = string.ascii_uppercase + string.digits
+CHAR2IDX = {c: i + 1 for i, c in enumerate(CHARS)}
+CHAR2IDX["<BLANK>"] = 0
+BLANK_IDX = 0
+IDX2CHAR = {v: k for k, v in CHAR2IDX.items()}
+IMAGE_HEIGHT = 32
+IMAGE_WIDTH = 128
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+font_path = None
+ocr_model = None
+
+
+# --------- Dataset --------- #
+class OCRDataset(Dataset):
+    def __init__(self, font_path, size=1000):
+        self.font = ImageFont.truetype(font_path, 32)
+        self.samples = ["".join(np.random.choice(list(CHARS), np.random.randint(4, 7)))
+                        for _ in range(size)]
+
+        self.transform = transforms.Compose([
+            transforms.Grayscale(),
+            transforms.Resize((IMAGE_HEIGHT, IMAGE_WIDTH)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.5,), (0.5,))
+        ])
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        text = self.samples[idx]
+        img = self.render_text(text)
+        img = self.transform(img)  # convert PIL to tensor with normalization
+
+        label = torch.tensor([CHAR2IDX[c] for c in text], dtype=torch.long)
+        return img, label
+
+
+    def render_text(self, text):
+        img = Image.new("L", (IMAGE_WIDTH, IMAGE_HEIGHT), color=255)
+        draw = ImageDraw.Draw(img)
+        bbox = self.font.getbbox(text)
+        w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+        draw.text(((IMAGE_WIDTH - w) // 2, (IMAGE_HEIGHT - h) // 2), text, font=self.font, fill=0)
+        return img
+
+
+# --------- Model --------- #
+class OCRModel(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d((2, 2), (2, 1)),  # height↓2, width↓1
+            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d((2, 2), (2, 1))  # height↓2 again, width↓1
+        )
+
+
+        self.rnn = nn.LSTM(64 * 8, 128, bidirectional=True, num_layers=2, batch_first=True)
+        self.fc = nn.Linear(256, num_classes)
+        with torch.no_grad():
+            self.fc.bias[0] = -5.0  # discourage blank early on
+
+
+    def forward(self, x):
+        b, c, h, w = x.size()
+        x = self.conv(x)
+        x = x.permute(0, 3, 1, 2)
+        x = x.reshape(b, x.size(1), -1)
+        x, _ = self.rnn(x)
+        x = self.fc(x)
+        return x
+
+
+def greedy_decode(log_probs):
+    # log_probs shape: (T, B, C)
+    # Usually, B=1 during inference
+    pred = log_probs.argmax(2).squeeze(1).tolist()  # this should give a list of ints
+    print(f"Decoded indices: {pred}")  # debug print
+    
+    decoded = []
+    prev = BLANK_IDX
+    for p in pred:
+        if p != prev and p != BLANK_IDX:
+            decoded.append(IDX2CHAR.get(p, ""))
+        prev = p
+    return ''.join(decoded)
+
+
+
+
+# --------- Custom Collate --------- #
+def custom_collate_fn(batch):
+    images, labels = zip(*batch)
+    images = torch.stack(images, 0)
+
+    flat_labels = []
+    label_lengths = []
+
+    for label in labels:
+        flat_labels.append(label)
+        label_lengths.append(len(label))
+
+    targets = torch.cat(flat_labels)
+    return images, targets, torch.tensor(label_lengths, dtype=torch.long)
+
+
+# --------- Model Save/Load --------- #
+def list_saved_models():
+    return [f for f in os.listdir() if f.endswith(".pth")]
+
+
+def save_model(model, path):
+    torch.save(model.state_dict(), path)
+
+
+def load_model(path):
+    global ocr_model
+    model = OCRModel(num_classes=len(CHAR2IDX))
+    model.load_state_dict(torch.load(path, map_location=device))
+    model.to(device)
+    model.eval()
+    ocr_model = model
+    return f"Model '{path}' loaded."
+
+
+# --------- Gradio Functions --------- #
+def train_model(font_file, epochs=100, learning_rate=0.001):
+    global font_path, ocr_model
+
+    # Save the uploaded font file
+    font_name = os.path.splitext(os.path.basename(font_file.name))[0]
+    font_path = f"./{font_name}.ttf"
+    with open(font_file.name, "rb") as uploaded:
+        with open(font_path, "wb") as f:
+            f.write(uploaded.read())
+
+    # Load dataset
+    dataset = OCRDataset(font_path)
+    dataloader = DataLoader(dataset, batch_size=16, shuffle=True, collate_fn=custom_collate_fn)
+
+    # Visualize one sample for sanity check
+    img, label = dataset[0]
+    print("Label:", ''.join([IDX2CHAR[i.item()] for i in label]))
+    plt.imshow(img.permute(1, 2, 0).squeeze(), cmap='gray')
+    plt.show()
+
+    # Initialize model
+    model = OCRModel(num_classes=len(CHAR2IDX)).to(device)
+    criterion = nn.CTCLoss(blank=0)
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+    # Training loop
+    for epoch in range(epochs):
+        for img, targets, target_lengths in dataloader:
+            img = img.to(device)
+            targets = targets.to(device)
+            target_lengths = target_lengths.to(device)
+
+            output = model(img)
+            batch_size = img.size(0)
+            seq_len = output.size(1)
+            input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)
+
+            loss = criterion(output.log_softmax(2).transpose(0, 1), targets, input_lengths, target_lengths)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        print(f"Epoch {epoch + 1}, Loss: {loss.item():.4f}")
+
+    # Save model with structured name
+    model_name = f"{font_name}_{epochs}epochs_lr{learning_rate:.0e}.pth"
+    save_model(model, model_name)
+    ocr_model = model
+    return f"Training complete! Model saved as '{model_name}'."
+
+
+
+
+def preprocess_image(image: Image.Image):
+    img_cv = np.array(image.convert("L"))
+
+    img_bin = cv2.adaptiveThreshold(img_cv, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+                                    cv2.THRESH_BINARY_INV, 25, 15)
+
+    # Invert if background is dark
+    white_px = (img_bin == 255).sum()
+    black_px = (img_bin == 0).sum()
+    if black_px > white_px:
+        img_bin = 255 - img_bin
+
+    # Resize and pad/crop to (IMAGE_HEIGHT, IMAGE_WIDTH)
+    h, w = img_bin.shape
+    scale = IMAGE_HEIGHT / h
+    new_w = int(w * scale)
+    resized = cv2.resize(img_bin, (new_w, IMAGE_HEIGHT), interpolation=cv2.INTER_AREA)
+
+    if new_w < IMAGE_WIDTH:
+        pad_width = IMAGE_WIDTH - new_w
+        padded = np.pad(resized, ((0, 0), (0, pad_width)), constant_values=255)
+    else:
+        padded = resized[:, :IMAGE_WIDTH]
+
+    return to_pil_image(padded)
+
+
+def predict_text(image: Image.Image):
+    if ocr_model is None:
+        return "Please load or train a model first."
+
+    processed = preprocess_image(image)
+
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))
+    ])
+    img_tensor = transform(processed).unsqueeze(0).to(device)  # (1, C, H, W)
+
+    with torch.no_grad():
+        output = ocr_model(img_tensor)  # (B, T, C)
+        log_probs = output.log_softmax(2).permute(1, 0, 2)  # (T, B, C)
+
+        pred = greedy_decode(log_probs)  # should be a string now
+
+        probs = log_probs.exp()
+        max_probs = probs.max(2)[0].squeeze(1)  # (T,)
+        avg_conf = max_probs.mean().item()
+
+    return f"Prediction: {pred}\nConfidence: {avg_conf:.2%}"
+
+
+
+
+# New helper function: generate label images grid
+def generate_labels(font_file=None, num_labels: int = 20):
+    global font_path
+
+    try:
+        if font_file:
+            font_path = "./temp_font_labels.ttf"
+            with open(font_file.name, "rb") as uploaded:
+                with open(font_path, "wb") as f:
+                    f.write(uploaded.read())
+
+        if font_path is None or not os.path.exists(font_path):
+            font = ImageFont.load_default()
+        else:
+            font = ImageFont.truetype(font_path, 32)
+
+        labels = ["".join(np.random.choice(list(CHARS), np.random.randint(4, 7))) for _ in range(num_labels)]
+
+        cols = min(5, num_labels)
+        rows = math.ceil(num_labels / cols)
+        cell_w, cell_h = IMAGE_WIDTH, IMAGE_HEIGHT
+
+        grid_img = Image.new("L", (cols * cell_w, rows * cell_h), color=255)
+        draw = ImageDraw.Draw(grid_img)
+
+        spacing = 0  # <-- spacing between characters
+
+        for idx, label in enumerate(labels):
+            x = (idx % cols) * cell_w
+            y = (idx // cols) * cell_h
+
+            # Draw each character with spacing
+            char_x = x + 10  # small left margin
+            char_y = y + (cell_h - font.size) // 2
+
+            for char in label:
+                draw.text((char_x, char_y), char, font=font, fill=0)
+                char_w = font.getbbox(char)[2] - font.getbbox(char)[0]
+                char_x += char_w + spacing  # move right with spacing
+
+        return grid_img
+
+    except Exception as e:
+        print("Error in generate_labels:", e)
+        error_img = Image.new("RGB", (512, 128), color=(255, 255, 255))
+        draw = ImageDraw.Draw(error_img)
+        draw.text((10, 50), f"Error: {str(e)}", fill=(255, 0, 0))
+        return error_img
+
+
+
+# --------- Updated Gradio UI with new tab --------- #
+with gr.Blocks() as demo:
+    with gr.Tab("1. Upload Font & Train"):
+        font_file = gr.File(label="Upload .ttf or .otf font", file_types=[".ttf", ".otf"])
+        epochs_input = gr.Slider(minimum=1, maximum=4096, value=256, step=1, label="Epochs")
+        lr_input = gr.Slider(minimum=0.001, maximum=0.1, value=0.05, step=0.001, label="Learning Rate")
+        train_button = gr.Button("Train OCR Model")
+        train_status = gr.Textbox(label="Status")
+
+    train_button.click(fn=train_model, inputs=[font_file, epochs_input, lr_input], outputs=train_status)
+
+
+    with gr.Tab("2. Use Trained Model"):
+        model_list = gr.Dropdown(choices=list_saved_models(), label="Select OCR Model")
+        refresh_btn = gr.Button("🔄 Refresh Models")
+        load_model_btn = gr.Button("Load Model")  # <-- new button
+
+        image_input = gr.Image(type="pil", label="Upload word strip")
+        predict_btn = gr.Button("Predict")
+        output_text = gr.Textbox(label="Recognized Text")
+        model_status = gr.Textbox(label="Model Load Status")
+
+        # Refresh dropdown choices
+        refresh_btn.click(fn=lambda: gr.update(choices=list_saved_models()), outputs=model_list)
+
+        # Load model on button click, NOT dropdown change
+        load_model_btn.click(fn=load_model, inputs=model_list, outputs=model_status)
+
+        predict_btn.click(fn=predict_text, inputs=image_input, outputs=output_text)
+
+    with gr.Tab("3. Generate Labels"):
+        font_file_labels = gr.File(label="Optional font for label image", file_types=[".ttf", ".otf"])
+        num_labels = gr.Number(value=20, label="Number of labels to generate", precision=0, interactive=True)
+        gen_button = gr.Button("Generate Label Grid")
+        label_image = gr.Image(label="Generated Labels", type="pil")
+
+        gen_button.click(fn=generate_labels, inputs=[font_file_labels, num_labels], outputs=label_image)
+
+
+
+if __name__ == "__main__":
+    demo.launch()