3morrrrr commited on
Commit
6fbeeae
·
verified ·
1 Parent(s): 9f070e0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +160 -0
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)