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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -50
app.py CHANGED
@@ -1,22 +1,24 @@
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
@@ -24,13 +26,41 @@ 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):
@@ -49,9 +79,18 @@ def detect_paper_angle(image, bounding_box):
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.")
@@ -78,39 +117,21 @@ def process_image(image, text):
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"
@@ -132,29 +153,29 @@ def gradio_inference(image, text):
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)
 
 
1
  import gradio as gr
2
  import logging
3
  from roboflow import Roboflow
4
+ from PIL import Image, ImageDraw
5
  import cv2
6
  import numpy as np
7
+ from selenium import webdriver
8
+ from selenium.webdriver.common.by import By
9
+ from selenium.webdriver.support.ui import WebDriverWait
10
+ from selenium.webdriver.support import expected_conditions as EC
11
+ from selenium.webdriver import ActionChains
12
+ from selenium.webdriver.support.ui import Select
13
+ import time
14
  import os
15
  from math import atan2, degrees
 
 
16
 
17
  # Configure logging
18
  logging.basicConfig(
19
  level=logging.DEBUG,
20
+ format="%(asctime)s - %(levelname)s - %(message)s",
21
+ handlers=[logging.FileHandler("debug.log"), logging.StreamHandler()],
 
 
 
22
  )
23
 
24
  # Roboflow and model configuration
 
26
  PROJECT_NAME = "model_verification_project"
27
  VERSION_NUMBER = 2
28
 
29
+ # Selenium configuration for Calligrapher
30
+ def get_calligrapher():
31
+ calli_url = "https://www.calligrapher.ai"
32
+ driver = webdriver.Chrome()
33
+ driver.maximize_window()
34
+ driver.get(calli_url)
35
+
36
+ # Adjust sliders for customization
37
+ speed_slider = WebDriverWait(driver, 20).until(EC.element_to_be_clickable((By.ID, 'speed-slider')))
38
+ ActionChains(driver).drag_and_drop_by_offset(speed_slider, 40, 0).perform()
39
+
40
+ bias_slider = WebDriverWait(driver, 20).until(EC.element_to_be_clickable((By.ID, 'bias-slider')))
41
+ ActionChains(driver).drag_and_drop_by_offset(bias_slider, 20, 0).perform()
42
+
43
+ width_slider = WebDriverWait(driver, 20).until(EC.element_to_be_clickable((By.ID, 'width-slider')))
44
+ ActionChains(driver).drag_and_drop_by_offset(width_slider, 20, 0).perform()
45
+
46
+ # Select handwriting style
47
+ select = Select(driver.find_element(By.ID, 'select-style'))
48
+ select.select_by_visible_text('9') # Adjust to the desired style
49
+ return driver
50
+
51
+ def get_calligrapher_text(driver, text, save_path):
52
+ text_input = WebDriverWait(driver, 20).until(EC.element_to_be_clickable((By.ID, 'text-input')))
53
+ text_input.clear()
54
+ text_input.send_keys(text)
55
+
56
+ draw_button = WebDriverWait(driver, 20).until(EC.element_to_be_clickable((By.ID, 'draw-button')))
57
+ draw_button.click()
58
+ time.sleep(3)
59
+
60
+ # Save the generated handwriting as an image
61
+ canvas = WebDriverWait(driver, 20).until(EC.presence_of_element_located((By.ID, 'draw-area')))
62
+ canvas.screenshot(save_path)
63
+ print(f"Handwriting saved to: {save_path}")
64
 
65
  # Function to detect paper angle within bounding box
66
  def detect_paper_angle(image, bounding_box):
 
79
  else:
80
  return 0
81
 
82
+ # Function to process image and overlay handwriting
83
  def process_image(image, text):
84
  try:
85
+ # Initialize Selenium and generate handwriting
86
+ save_path = "/tmp/handwriting.png"
87
+ driver = get_calligrapher()
88
+ get_calligrapher_text(driver, text, save_path)
89
+ driver.quit()
90
+
91
+ # Open generated handwriting image
92
+ handwriting_image = Image.open(save_path).convert("RGBA")
93
+
94
  # Initialize Roboflow
95
  rf = Roboflow(api_key=ROBOFLOW_API_KEY)
96
  logging.debug("Initialized Roboflow API.")
 
117
  for obj in prediction['predictions']:
118
  white_paper_width = obj['width']
119
  white_paper_height = obj['height']
120
+
121
+ padding_x = int(white_paper_width * 0.1) # 10% padding horizontally
122
+ padding_y = int(white_paper_height * 0.1) # 10% padding vertically
123
+
124
  box_width = white_paper_width - 2 * padding_x
125
  box_height = white_paper_height - 2 * padding_y
 
126
 
127
  x1_padded = int(obj['x'] - white_paper_width / 2 + padding_x)
128
  y1_padded = int(obj['y'] - white_paper_height / 2 + padding_y)
 
 
 
 
 
 
 
 
 
 
 
129
 
130
+ # Resize handwriting image to fit the detected area
131
+ resized_handwriting = handwriting_image.resize((box_width, box_height), Image.ANTIALIAS)
132
 
133
+ # Paste handwriting onto detected area
134
+ pil_image.paste(resized_handwriting, (x1_padded, y1_padded), resized_handwriting)
 
 
 
 
 
 
 
 
135
 
136
  # Save and return output image path
137
  output_image_path = "/tmp/output_image.png"
 
153
  logging.error("Gradio inference failed.")
154
  return None, None, "An error occurred while processing the image. Please check the logs."
155
 
 
156
  # Gradio interface
157
  interface = gr.Interface(
158
  fn=gradio_inference,
159
  inputs=[
160
+ gr.Image(type="pil", label="Upload an Image"),
161
+ gr.Textbox(label="Enter Text to Overlay"),
162
  ],
163
  outputs=[
164
+ gr.Image(label="Processed Image Preview"), # Preview processed image
165
  gr.File(label="Download Processed Image"), # Download the image
166
  gr.Textbox(label="Status"), # Status message
167
  ],
168
+ title="Roboflow Detection with Calligrapher Text Overlay",
169
  description=(
170
+ "Upload an image, enter text to overlay, and let the Roboflow model process the image. "
171
+ "Handwritten text is generated using Calligrapher.ai and overlaid on the detected white paper areas. "
172
  "Preview or download the output image below."
173
  ),
174
+ allow_flagging="never",
175
  )
176
 
177
  # Launch the Gradio app
178
  if __name__ == "__main__":
179
  logging.debug("Launching Gradio interface.")
180
  interface.launch(share=True)
181
+