import gradio as gr import logging from roboflow import Roboflow from PIL import Image, ImageDraw, ImageFont, ImageFilter import cv2 import numpy as np import os from math import atan2, degrees from diffusers import AutoPipelineForText2Image import torch # Configure logging logging.basicConfig( level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.FileHandler("debug.log"), logging.StreamHandler() ] ) # Roboflow and model configuration ROBOFLOW_API_KEY = "KUP9w62eUcD5PrrRMJsV" # Replace with your API key PROJECT_NAME = "model_verification_project" VERSION_NUMBER = 2 # Initialize the FLUX handwriting model device = "cuda" if torch.cuda.is_available() else "cpu" pipeline = AutoPipelineForText2Image.from_pretrained( 'black-forest-labs/FLUX.1-dev', torch_dtype=torch.float16 ).to(device) pipeline.load_lora_weights('fofr/flux-handwriting', weight_name='lora.safetensors') # Function to detect paper angle within bounding box def detect_paper_angle(image, bounding_box): x1, y1, x2, y2 = bounding_box roi = np.array(image)[y1:y2, x1:x2] gray = cv2.cvtColor(roi, cv2.COLOR_RGBA2GRAY) edges = cv2.Canny(gray, 50, 150) lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=50, maxLineGap=10) if lines is not None: longest_line = max(lines, key=lambda line: np.linalg.norm((line[0][2] - line[0][0], line[0][3] - line[0][1]))) x1, y1, x2, y2 = longest_line[0] dx = x2 - x1 dy = y2 - y1 angle = degrees(atan2(dy, dx)) return angle else: return 0 # Function to process image and overlay text def process_image(image, text): try: # Initialize Roboflow rf = Roboflow(api_key=ROBOFLOW_API_KEY) logging.debug("Initialized Roboflow API.") project = rf.workspace().project(PROJECT_NAME) logging.debug("Accessed project in Roboflow.") model = project.version(VERSION_NUMBER).model logging.debug("Loaded model from Roboflow.") # Save input image temporarily input_image_path = "/tmp/input_image.jpg" image.save(input_image_path) logging.debug(f"Input image saved to {input_image_path}.") # Perform inference logging.debug("Performing inference on the image...") prediction = model.predict(input_image_path, confidence=70, overlap=50).json() logging.debug(f"Inference result: {prediction}") # Open the image for processing pil_image = image.convert("RGBA") logging.debug("Converted image to RGBA mode.") # Iterate over detected objects for obj in prediction['predictions']: white_paper_width = obj['width'] white_paper_height = obj['height'] padding_x = int(white_paper_width * 0.1) padding_y = int(white_paper_height * 0.1) box_width = white_paper_width - 2 * padding_x box_height = white_paper_height - 2 * padding_y logging.debug(f"Padded white paper dimensions: width={box_width}, height={box_height}.") x1_padded = int(obj['x'] - white_paper_width / 2 + padding_x) y1_padded = int(obj['y'] - white_paper_height / 2 + padding_y) x2_padded = int(obj['x'] + white_paper_width / 2 - padding_x) y2_padded = int(obj['y'] + white_paper_height / 2 - padding_y) # Detect paper angle angle = detect_paper_angle(np.array(image), (x1_padded, y1_padded, x2_padded, y2_padded)) logging.debug(f"Detected paper angle: {angle} degrees.") # Generate handwriting image with transparent background prompt = f'HWRIT handwriting saying "{text}", neat style, black ink on transparent background' generated_image = pipeline(prompt).images[0].convert("RGBA") logging.debug("Generated handwriting image.") # Resize generated handwriting to fit the detected area generated_image = generated_image.resize((box_width, box_height), Image.ANTIALIAS) # Create a mask for the generated handwriting mask = generated_image.split()[3] # Rotate the generated handwriting to match the detected paper angle rotated_handwriting = generated_image.rotate(-angle, resample=Image.BICUBIC, center=(box_width // 2, box_height // 2)) mask = mask.rotate(-angle, resample=Image.BICUBIC, center=(box_width // 2, box_height // 2)) # Paste the rotated handwriting onto the original image pil_image.paste(rotated_handwriting, (x1_padded, y1_padded), mask) logging.debug("Pasted generated handwriting onto the original image.") # Save and return output image path output_image_path = "/tmp/output_image.png" pil_image.convert("RGB").save(output_image_path) logging.debug(f"Output image saved to {output_image_path}.") return output_image_path except Exception as e: logging.error(f"Error during image processing: {e}") return None # Gradio interface function def gradio_inference(image, text): logging.debug("Starting Gradio inference.") result_path = process_image(image, text) if result_path: logging.debug("Gradio inference successful.") return result_path, result_path, "Processing complete! Download the image below." logging.error("Gradio inference failed.") return None, None, "An error occurred while processing the image. Please check the logs." # Gradio interface # Gradio interface interface = gr.Interface( fn=gradio_inference, inputs=[ gr.Image(type="pil", label="Upload an Image"), # Upload an image gr.Textbox(label="Enter Text to Overlay"), # Enter text to overlay ], outputs=[ gr.Image(label="Processed Image Preview"), # Preview the processed image gr.File(label="Download Processed Image"), # Download the image gr.Textbox(label="Status"), # Status message ], title="Handwriting Overlay on White Paper", description=( "Upload an image with white paper detected, and enter the text to overlay. " "This app will generate handwriting using the FLUX handwriting model and overlay it on the detected white paper. " "Preview or download the output image below." ), allow_flagging="never", # Disables flagging ) # Launch the Gradio app if __name__ == "__main__": logging.debug("Launching Gradio interface.") interface.launch(share=True)