Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,160 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import logging
|
3 |
+
from roboflow import Roboflow
|
4 |
+
from PIL import Image, ImageDraw, ImageFont, ImageFilter
|
5 |
+
import cv2
|
6 |
+
import numpy as np
|
7 |
+
import os
|
8 |
+
from math import atan2, degrees
|
9 |
+
from diffusers import AutoPipelineForText2Image
|
10 |
+
import torch
|
11 |
+
|
12 |
+
# Configure logging
|
13 |
+
logging.basicConfig(
|
14 |
+
level=logging.DEBUG,
|
15 |
+
format='%(asctime)s - %(levelname)s - %(message)s',
|
16 |
+
handlers=[
|
17 |
+
logging.FileHandler("debug.log"),
|
18 |
+
logging.StreamHandler()
|
19 |
+
]
|
20 |
+
)
|
21 |
+
|
22 |
+
# Roboflow and model configuration
|
23 |
+
ROBOFLOW_API_KEY = "KUP9w62eUcD5PrrRMJsV" # Replace with your API key
|
24 |
+
PROJECT_NAME = "model_verification_project"
|
25 |
+
VERSION_NUMBER = 2
|
26 |
+
|
27 |
+
# Initialize the FLUX handwriting model
|
28 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
29 |
+
pipeline = AutoPipelineForText2Image.from_pretrained(
|
30 |
+
'black-forest-labs/FLUX.1-dev',
|
31 |
+
torch_dtype=torch.float16
|
32 |
+
).to(device)
|
33 |
+
pipeline.load_lora_weights('fofr/flux-handwriting', weight_name='lora.safetensors')
|
34 |
+
|
35 |
+
# Function to detect paper angle within bounding box
|
36 |
+
def detect_paper_angle(image, bounding_box):
|
37 |
+
x1, y1, x2, y2 = bounding_box
|
38 |
+
roi = np.array(image)[y1:y2, x1:x2]
|
39 |
+
gray = cv2.cvtColor(roi, cv2.COLOR_RGBA2GRAY)
|
40 |
+
edges = cv2.Canny(gray, 50, 150)
|
41 |
+
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=50, maxLineGap=10)
|
42 |
+
if lines is not None:
|
43 |
+
longest_line = max(lines, key=lambda line: np.linalg.norm((line[0][2] - line[0][0], line[0][3] - line[0][1])))
|
44 |
+
x1, y1, x2, y2 = longest_line[0]
|
45 |
+
dx = x2 - x1
|
46 |
+
dy = y2 - y1
|
47 |
+
angle = degrees(atan2(dy, dx))
|
48 |
+
return angle
|
49 |
+
else:
|
50 |
+
return 0
|
51 |
+
|
52 |
+
# Function to process image and overlay text
|
53 |
+
def process_image(image, text):
|
54 |
+
try:
|
55 |
+
# Initialize Roboflow
|
56 |
+
rf = Roboflow(api_key=ROBOFLOW_API_KEY)
|
57 |
+
logging.debug("Initialized Roboflow API.")
|
58 |
+
project = rf.workspace().project(PROJECT_NAME)
|
59 |
+
logging.debug("Accessed project in Roboflow.")
|
60 |
+
model = project.version(VERSION_NUMBER).model
|
61 |
+
logging.debug("Loaded model from Roboflow.")
|
62 |
+
|
63 |
+
# Save input image temporarily
|
64 |
+
input_image_path = "/tmp/input_image.jpg"
|
65 |
+
image.save(input_image_path)
|
66 |
+
logging.debug(f"Input image saved to {input_image_path}.")
|
67 |
+
|
68 |
+
# Perform inference
|
69 |
+
logging.debug("Performing inference on the image...")
|
70 |
+
prediction = model.predict(input_image_path, confidence=70, overlap=50).json()
|
71 |
+
logging.debug(f"Inference result: {prediction}")
|
72 |
+
|
73 |
+
# Open the image for processing
|
74 |
+
pil_image = image.convert("RGBA")
|
75 |
+
logging.debug("Converted image to RGBA mode.")
|
76 |
+
|
77 |
+
# Iterate over detected objects
|
78 |
+
for obj in prediction['predictions']:
|
79 |
+
white_paper_width = obj['width']
|
80 |
+
white_paper_height = obj['height']
|
81 |
+
padding_x = int(white_paper_width * 0.1)
|
82 |
+
padding_y = int(white_paper_height * 0.1)
|
83 |
+
box_width = white_paper_width - 2 * padding_x
|
84 |
+
box_height = white_paper_height - 2 * padding_y
|
85 |
+
logging.debug(f"Padded white paper dimensions: width={box_width}, height={box_height}.")
|
86 |
+
|
87 |
+
x1_padded = int(obj['x'] - white_paper_width / 2 + padding_x)
|
88 |
+
y1_padded = int(obj['y'] - white_paper_height / 2 + padding_y)
|
89 |
+
x2_padded = int(obj['x'] + white_paper_width / 2 - padding_x)
|
90 |
+
y2_padded = int(obj['y'] + white_paper_height / 2 - padding_y)
|
91 |
+
|
92 |
+
# Detect paper angle
|
93 |
+
angle = detect_paper_angle(np.array(image), (x1_padded, y1_padded, x2_padded, y2_padded))
|
94 |
+
logging.debug(f"Detected paper angle: {angle} degrees.")
|
95 |
+
|
96 |
+
# Generate handwriting image with transparent background
|
97 |
+
prompt = f'HWRIT handwriting saying "{text}", neat style, black ink on transparent background'
|
98 |
+
generated_image = pipeline(prompt).images[0].convert("RGBA")
|
99 |
+
logging.debug("Generated handwriting image.")
|
100 |
+
|
101 |
+
# Resize generated handwriting to fit the detected area
|
102 |
+
generated_image = generated_image.resize((box_width, box_height), Image.ANTIALIAS)
|
103 |
+
|
104 |
+
# Create a mask for the generated handwriting
|
105 |
+
mask = generated_image.split()[3]
|
106 |
+
|
107 |
+
# Rotate the generated handwriting to match the detected paper angle
|
108 |
+
rotated_handwriting = generated_image.rotate(-angle, resample=Image.BICUBIC, center=(box_width // 2, box_height // 2))
|
109 |
+
mask = mask.rotate(-angle, resample=Image.BICUBIC, center=(box_width // 2, box_height // 2))
|
110 |
+
|
111 |
+
# Paste the rotated handwriting onto the original image
|
112 |
+
pil_image.paste(rotated_handwriting, (x1_padded, y1_padded), mask)
|
113 |
+
logging.debug("Pasted generated handwriting onto the original image.")
|
114 |
+
|
115 |
+
# Save and return output image path
|
116 |
+
output_image_path = "/tmp/output_image.png"
|
117 |
+
pil_image.convert("RGB").save(output_image_path)
|
118 |
+
logging.debug(f"Output image saved to {output_image_path}.")
|
119 |
+
return output_image_path
|
120 |
+
|
121 |
+
except Exception as e:
|
122 |
+
logging.error(f"Error during image processing: {e}")
|
123 |
+
return None
|
124 |
+
|
125 |
+
# Gradio interface function
|
126 |
+
def gradio_inference(image, text):
|
127 |
+
logging.debug("Starting Gradio inference.")
|
128 |
+
result_path = process_image(image, text)
|
129 |
+
if result_path:
|
130 |
+
logging.debug("Gradio inference successful.")
|
131 |
+
return result_path, result_path, "Processing complete! Download the image below."
|
132 |
+
logging.error("Gradio inference failed.")
|
133 |
+
return None, None, "An error occurred while processing the image. Please check the logs."
|
134 |
+
|
135 |
+
# Gradio interface
|
136 |
+
# Gradio interface
|
137 |
+
interface = gr.Interface(
|
138 |
+
fn=gradio_inference,
|
139 |
+
inputs=[
|
140 |
+
gr.Image(type="pil", label="Upload an Image"), # Upload an image
|
141 |
+
gr.Textbox(label="Enter Text to Overlay"), # Enter text to overlay
|
142 |
+
],
|
143 |
+
outputs=[
|
144 |
+
gr.Image(label="Processed Image Preview"), # Preview the processed image
|
145 |
+
gr.File(label="Download Processed Image"), # Download the image
|
146 |
+
gr.Textbox(label="Status"), # Status message
|
147 |
+
],
|
148 |
+
title="Handwriting Overlay on White Paper",
|
149 |
+
description=(
|
150 |
+
"Upload an image with white paper detected, and enter the text to overlay. "
|
151 |
+
"This app will generate handwriting using the FLUX handwriting model and overlay it on the detected white paper. "
|
152 |
+
"Preview or download the output image below."
|
153 |
+
),
|
154 |
+
allow_flagging="never", # Disables flagging
|
155 |
+
)
|
156 |
+
|
157 |
+
# Launch the Gradio app
|
158 |
+
if __name__ == "__main__":
|
159 |
+
logging.debug("Launching Gradio interface.")
|
160 |
+
interface.launch(share=True)
|