app.py

import gradio as gr
from inference import OneDMInference
import os
from PIL import Image
import cv2
import numpy as np
import torch
import torch.nn.functional as F

# Load the model
model = OneDMInference(
    model_path='one_dm_finetuned.pt',
    cfg_path='configs/finetuned.yml'
)

# Define Laplacian kernel (ensure it’s on the correct device if needed)
laplace = torch.tensor(
    [[0, 1, 0],
     [1, -4, 1],
     [0, 1, 0]], dtype=torch.float, requires_grad=False
).view(1, 1, 3, 3)

def generate_laplace_image(image_path, target_size=(64, 64)):
    """
    Generate a Laplace image from the input image using a Laplacian filter.
    Adjusted to match model-expected dimensions (e.g., 64x64).
    """
    # Read image
    image = cv2.imread(image_path)
    if image is None:
        raise ValueError(f"Could not read image at {image_path}")

    # Convert to grayscale
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Resize to model-compatible size (e.g., 64x64)
    image = cv2.resize(image, target_size)

    # Convert to tensor
    x = torch.from_numpy(image).unsqueeze(0).unsqueeze(0).float()

    # Normalize input
    x = x / 255.0

    # Apply Laplacian filter with proper padding
    y = F.conv2d(x, laplace, stride=1, padding=1)  # Padding=1 keeps spatial dims intact

    # Process output
    y = y.squeeze().numpy()
    y = np.clip(y * 255.0, 0, 255)
    y = y.astype(np.uint8)

    # Apply thresholding
    _, threshold = cv2.threshold(y, 0, 255, cv2.THRESH_OTSU)

    # Save output
    laplace_path = os.path.splitext(image_path)[0] + "_laplace.png"
    cv2.imwrite(laplace_path, threshold)

    return laplace_path
def generate_handwriting(text, style_image, laplace_image=None):
    output_dir = "./generated"
    os.makedirs(output_dir, exist_ok=True)

    # Assume model expects 64x64 inputs based on logs (adjust if config specifies otherwise)
    target_size = (64, 64)

    # Generate Laplace image if not provided
    if laplace_image is None:
        laplace_image = generate_laplace_image(style_image, target_size)
    else:
        # Ensure provided Laplace image matches expected size
        laplace_img = cv2.imread(laplace_image, cv2.IMREAD_GRAYSCALE)
        if laplace_img.shape != target_size:
            laplace_img = cv2.resize(laplace_img, target_size)
            laplace_image = os.path.splitext(laplace_image)[0] + "_resized.png"
            cv2.imwrite(laplace_image, laplace_img)

    # Resize style image to match model expectations
    style_img = cv2.imread(style_image)
    style_img_resized = cv2.resize(style_img, target_size)
    style_image_resized = os.path.splitext(style_image)[0] + "_resized.png"
    cv2.imwrite(style_image_resized, style_img_resized)

    # Generate handwriting for each word
    words = text.split()
    generated_image_paths = []
    for word in words:
        output_paths = model.generate(
            text=word,
            style_path=style_image_resized,  # Use resized style image
            laplace_path=laplace_image,      # Use Laplace image
            output_dir=output_dir
        )
        generated_image_paths.append(output_paths[0])

    # Load generated images
    images = [Image.open(img_path) for img_path in generated_image_paths]

    # Constants for spacing and margins (adjusted for better spacing)
    word_gap = 5  # Reduced from 20 to 5 for closer word spacing
    line_gap = 20  # Reduced from 30 for tighter lines
    max_words_per_line = 5
    top_margin = 10  # Reduced from 30
    left_margin = 10  # Reduced from 30

    # Calculate line dimensions
    lines = []
    current_line = []
    current_line_width = 0
    current_line_height = 0

    for img in images:
        if len(current_line) >= max_words_per_line or current_line_width + img.size[0] > 500:  # Add a max width constraint (e.g., 500px)
            lines.append((current_line, current_line_width - word_gap, current_line_height))
            current_line = []
            current_line_width = 0
            current_line_height = 0

        current_line.append(img)
        current_line_width += img.size[0] + word_gap
        current_line_height = max(current_line_height, img.size[1])

    # Add the last line if it has content
    if current_line:
        lines.append((current_line, current_line_width - word_gap, current_line_height))

    # Calculate total dimensions
    total_width = max(line[1] for line in lines) + (2 * left_margin)  # Width of the widest line
    total_height = sum(line[2] for line in lines) + (len(lines) - 1) * line_gap + top_margin

    # Create merged image
    merged_image = Image.new('RGB', (total_width, total_height), color=(255, 255, 255))

    # Paste words into the image
    y_offset = top_margin
    for line_images, line_width, line_height in lines:
        x_offset = left_margin  # Align to the left instead of centering
        for img in line_images:
            # Adjust y_offset for each word to align baselines (optional, if heights vary significantly)
            word_y_offset = y_offset + (line_height - img.size[1])  # Align to the bottom of the line
            merged_image.paste(img, (x_offset, word_y_offset))
            x_offset += img.size[0] + word_gap
        y_offset += line_height + line_gap

    # Save merged image
    merged_image_path = os.path.join(output_dir, "merged_output.png")
    merged_image.save(merged_image_path)

    return merged_image_path, gr.update(value=laplace_image)


# Create Gradio interface
iface = gr.Interface(
    fn=generate_handwriting,
    inputs=[
        gr.Textbox(label="Text to generate"),
        gr.Image(label="Style Image", type="filepath"),
        gr.Image(label="Laplace Image (Optional)", type="filepath")
    ],
    outputs=[
        gr.Image(label="Generated Handwriting"),
        gr.Image(label="Laplace Image (Optional)")
    ],
    title="Handwriting Generation",
    description="Generate handwritten text using One-DM model. If no Laplace image is provided, it will be generated from the style image.",
    examples=[
        ["Hello World",
         "English_data/Dataset/test/169/c04-134-05-08.png",
         "English_data/Dataset_laplace/test/169/c04-134-00-00.png"]
    ]
)

if __name__ == "__main__":
    iface.launch(share=True)