Spaces:
Runtime error
Runtime error
Pratap2002
commited on
Commit
•
3ea6f67
1
Parent(s):
85b5c69
Upload 2 files
Browse files- app.py +188 -0
- requirements.txt +8 -0
app.py
ADDED
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import streamlit as st
|
3 |
+
import requests
|
4 |
+
from PIL import Image, ImageDraw, ImageFont
|
5 |
+
import io
|
6 |
+
import base64
|
7 |
+
import easyocr
|
8 |
+
import numpy as np
|
9 |
+
import cv2
|
10 |
+
|
11 |
+
# Set up logging
|
12 |
+
import logging
|
13 |
+
logging.basicConfig(level=logging.DEBUG)
|
14 |
+
logger = logging.getLogger(__name__)
|
15 |
+
|
16 |
+
# Hugging Face API setup
|
17 |
+
API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
|
18 |
+
|
19 |
+
headers = {"Authorization": f"Bearer {HF_TOKEN}"}
|
20 |
+
|
21 |
+
# Initialize EasyOCR reader
|
22 |
+
reader = easyocr.Reader(['en'])
|
23 |
+
|
24 |
+
def query(payload):
|
25 |
+
try:
|
26 |
+
response = requests.post(API_URL, headers=headers, json=payload)
|
27 |
+
response.raise_for_status()
|
28 |
+
|
29 |
+
logger.debug(f"API response status code: {response.status_code}")
|
30 |
+
logger.debug(f"API response headers: {response.headers}")
|
31 |
+
|
32 |
+
content_type = response.headers.get('Content-Type', '')
|
33 |
+
if 'application/json' in content_type:
|
34 |
+
return response.json()
|
35 |
+
elif 'image' in content_type:
|
36 |
+
return response.content
|
37 |
+
else:
|
38 |
+
logger.error(f"Unexpected content type: {content_type}")
|
39 |
+
st.error(f"Unexpected content type: {content_type}")
|
40 |
+
return None
|
41 |
+
except requests.exceptions.RequestException as e:
|
42 |
+
logger.error(f"Request failed: {str(e)}")
|
43 |
+
st.error(f"Request failed: {str(e)}")
|
44 |
+
return None
|
45 |
+
|
46 |
+
def increase_image_quality(image, scale_factor):
|
47 |
+
width, height = image.size
|
48 |
+
new_size = (width * scale_factor, height * scale_factor)
|
49 |
+
return image.resize(new_size, Image.LANCZOS)
|
50 |
+
|
51 |
+
def extract_text_from_image(image):
|
52 |
+
img_array = np.array(image)
|
53 |
+
results = reader.readtext(img_array)
|
54 |
+
return ' '.join([result[1] for result in results])
|
55 |
+
|
56 |
+
def remove_text_from_image(image, text_to_remove):
|
57 |
+
img_array = np.array(image)
|
58 |
+
results = reader.readtext(img_array)
|
59 |
+
|
60 |
+
for (bbox, text, prob) in results:
|
61 |
+
if text_to_remove.lower() in text.lower():
|
62 |
+
top_left = tuple(map(int, bbox[0]))
|
63 |
+
bottom_right = tuple(map(int, bbox[2]))
|
64 |
+
|
65 |
+
# Convert image to OpenCV format
|
66 |
+
img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
|
67 |
+
|
68 |
+
# Create a mask for inpainting
|
69 |
+
mask = np.zeros(img_cv.shape[:2], dtype=np.uint8)
|
70 |
+
cv2.rectangle(mask, top_left, bottom_right, (255, 255, 255), -1)
|
71 |
+
|
72 |
+
# Perform inpainting
|
73 |
+
inpainted = cv2.inpaint(img_cv, mask, 3, cv2.INPAINT_TELEA)
|
74 |
+
|
75 |
+
# Convert back to PIL Image
|
76 |
+
image = Image.fromarray(cv2.cvtColor(inpainted, cv2.COLOR_BGR2RGB))
|
77 |
+
|
78 |
+
return image, top_left, (bottom_right[0] - top_left[0], bottom_right[1] - top_left[1])
|
79 |
+
|
80 |
+
logger.warning(f"Text '{text_to_remove}' not found in the image.")
|
81 |
+
return image, None, None
|
82 |
+
|
83 |
+
def add_text_to_image(image, text, font_size=40, font_color="#FFFFFF", position=None, size=None):
|
84 |
+
draw = ImageDraw.Draw(image)
|
85 |
+
try:
|
86 |
+
font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
|
87 |
+
except IOError:
|
88 |
+
logger.warning("Roboto-Bold font not found, using default font")
|
89 |
+
font = ImageFont.load_default()
|
90 |
+
|
91 |
+
img_width, img_height = image.size
|
92 |
+
if position is None or size is None:
|
93 |
+
# Calculate the center position if no position is provided
|
94 |
+
bbox = font.getbbox(text)
|
95 |
+
text_width = bbox[2] - bbox[0]
|
96 |
+
text_height = bbox[3] - bbox[1]
|
97 |
+
position = ((img_width - text_width) // 2, (img_height - text_height) // 2)
|
98 |
+
size = (text_width, text_height)
|
99 |
+
|
100 |
+
# Adjust font size to fit within the given size
|
101 |
+
while font.getbbox(text)[2] - font.getbbox(text)[0] > size[0] or font.getbbox(text)[3] - font.getbbox(text)[1] > size[1]:
|
102 |
+
font_size -= 1
|
103 |
+
font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
|
104 |
+
|
105 |
+
# Use the exact position of the removed text
|
106 |
+
logger.debug(f"Adding text at position: {position}")
|
107 |
+
draw.text(position, text, font=font, fill=font_color)
|
108 |
+
return image
|
109 |
+
|
110 |
+
def main():
|
111 |
+
st.title("Poster Generator and Editor")
|
112 |
+
|
113 |
+
# Image Generation
|
114 |
+
st.header("Generate Poster")
|
115 |
+
poster_type = st.selectbox("Poster Type", ["Fashion", "Movie", "Event", "Advertisement", "Other"])
|
116 |
+
prompt = st.text_area("Prompt")
|
117 |
+
num_images = st.number_input("Number of Images", min_value=1, max_value=5, value=1)
|
118 |
+
quality_factor = st.number_input("Quality Factor", min_value=1, max_value=4, value=1)
|
119 |
+
|
120 |
+
if st.button("Generate Images"):
|
121 |
+
if poster_type == "Other":
|
122 |
+
full_prompt = f"A colorful poster with the following elements: {prompt}"
|
123 |
+
else:
|
124 |
+
full_prompt = f"A colorful {poster_type.lower()} poster with the following elements: {prompt}"
|
125 |
+
|
126 |
+
generated_images = []
|
127 |
+
for i in range(num_images):
|
128 |
+
with st.spinner(f"Generating image {i+1}..."):
|
129 |
+
logger.info(f"Generating image {i+1} with prompt: {full_prompt}")
|
130 |
+
response = query({"inputs": full_prompt})
|
131 |
+
|
132 |
+
if isinstance(response, bytes):
|
133 |
+
image = Image.open(io.BytesIO(response))
|
134 |
+
if quality_factor > 1:
|
135 |
+
image = increase_image_quality(image, quality_factor)
|
136 |
+
generated_images.append(image)
|
137 |
+
else:
|
138 |
+
st.error("Failed to generate image")
|
139 |
+
|
140 |
+
# Display generated images
|
141 |
+
for i, img in enumerate(generated_images):
|
142 |
+
st.image(img, caption=f"Generated Poster {i+1}", use_column_width=True)
|
143 |
+
|
144 |
+
# Save image to session state for editing
|
145 |
+
img_byte_arr = io.BytesIO()
|
146 |
+
img.save(img_byte_arr, format='PNG')
|
147 |
+
img_byte_arr = img_byte_arr.getvalue()
|
148 |
+
st.session_state[f'image_{i}'] = img_byte_arr
|
149 |
+
|
150 |
+
# Image Editing
|
151 |
+
st.header("Edit Poster")
|
152 |
+
image_to_edit = st.selectbox("Select Image to Edit", [f"Generated Poster {i+1}" for i in range(len(st.session_state.keys()))])
|
153 |
+
|
154 |
+
if image_to_edit:
|
155 |
+
image_index = int(image_to_edit.split()[-1]) - 1
|
156 |
+
img_bytes = st.session_state[f'image_{image_index}']
|
157 |
+
img = Image.open(io.BytesIO(img_bytes))
|
158 |
+
st.image(img, caption="Current Image", use_column_width=True)
|
159 |
+
|
160 |
+
text_to_remove = st.text_input("Text to Remove")
|
161 |
+
new_text = st.text_input("New Text")
|
162 |
+
font_size = st.number_input("Font Size", min_value=1, max_value=100, value=40)
|
163 |
+
font_color = st.color_picker("Font Color", "#FFFFFF")
|
164 |
+
|
165 |
+
if st.button("Apply Changes"):
|
166 |
+
position = None
|
167 |
+
size = None
|
168 |
+
if text_to_remove:
|
169 |
+
img, position, size = remove_text_from_image(img, text_to_remove)
|
170 |
+
|
171 |
+
if new_text:
|
172 |
+
img = add_text_to_image(img, new_text, font_size, font_color, position, size)
|
173 |
+
|
174 |
+
st.image(img, caption="Edited Image", use_column_width=True)
|
175 |
+
|
176 |
+
# Save edited image for download
|
177 |
+
img_byte_arr = io.BytesIO()
|
178 |
+
img.save(img_byte_arr, format='PNG')
|
179 |
+
img_byte_arr = img_byte_arr.getvalue()
|
180 |
+
st.download_button(
|
181 |
+
label="Download Edited Image",
|
182 |
+
data=img_byte_arr,
|
183 |
+
file_name="edited_poster.png",
|
184 |
+
mime="image/png"
|
185 |
+
)
|
186 |
+
|
187 |
+
if __name__ == "__main__":
|
188 |
+
main()
|
requirements.txt
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
streamlit
|
2 |
+
python-dotenv
|
3 |
+
requests
|
4 |
+
Pillow
|
5 |
+
easyocr
|
6 |
+
numpy
|
7 |
+
protobuf
|
8 |
+
altair
|