Spaces:
Running
Running
| import streamlit as st | |
| from PIL import Image, ImageDraw, ImageFont | |
| import io | |
| import cv2 | |
| import numpy as np | |
| import easyocr | |
| import os | |
| from dotenv import load_dotenv | |
| import requests | |
| import logging | |
| # Load environment variables | |
| load_dotenv() | |
| # Set up logging | |
| logging.basicConfig(level=logging.DEBUG) | |
| logger = logging.getLogger(__name__) | |
| # Hugging Face API setup | |
| API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell" | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| headers = {"Authorization": f"Bearer {HF_TOKEN}"} | |
| def load_image(image_file): | |
| img = Image.open(image_file) | |
| return img | |
| def detect_text(image): | |
| reader = easyocr.Reader(['en']) | |
| img_array = np.array(image) | |
| results = reader.readtext(img_array) | |
| return [(text, box) for (box, text, _) in results] | |
| def replace_text_in_image(image, text_to_replace, new_text): | |
| img_array = np.array(image) | |
| for (text, box) in detect_text(image): | |
| if text == text_to_replace: | |
| x, y, w, h = int(box[0][0]), int(box[0][1]), int(box[2][0] - box[0][0]), int(box[2][1] - box[0][1]) | |
| mask = np.zeros(img_array.shape[:2], dtype=np.uint8) | |
| cv2.rectangle(mask, (x, y), (x+w, y+h), 255, -1) | |
| img_array = cv2.inpaint(img_array, mask, 3, cv2.INPAINT_TELEA) | |
| image = Image.fromarray(img_array) | |
| draw = ImageDraw.Draw(image) | |
| font = ImageFont.truetype("arial.ttf", 40) | |
| draw.text((x, y), new_text, font=font, fill="#000000") | |
| return image | |
| return Image.fromarray(img_array) | |
| def query(payload): | |
| try: | |
| response = requests.post(API_URL, headers=headers, json=payload) | |
| response.raise_for_status() | |
| logger.debug(f"API response status code: {response.status_code}") | |
| logger.debug(f"API response headers: {response.headers}") | |
| content_type = response.headers.get('Content-Type', '') | |
| if 'application/json' in content_type: | |
| return response.json() | |
| elif 'image' in content_type: | |
| return response.content | |
| else: | |
| logger.error(f"Unexpected content type: {content_type}") | |
| st.error(f"Unexpected content type: {content_type}") | |
| return None | |
| except requests.exceptions.RequestException as e: | |
| logger.error(f"Request failed: {str(e)}") | |
| st.error(f"Request failed: {str(e)}") | |
| return None | |
| def increase_image_quality(image, scale_factor): | |
| width, height = image.size | |
| new_size = (width * scale_factor, height * scale_factor) | |
| return image.resize(new_size, Image.LANCZOS) | |
| def image_text_replacer(): | |
| st.header("Image Text Replacer") | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| image = load_image(uploaded_file) | |
| st.image(image, caption='Uploaded Image', use_column_width=True) | |
| text_results = detect_text(image) | |
| st.subheader("Detected Text:") | |
| edited_image = image.copy() | |
| for i, (text, box) in enumerate(text_results): | |
| if text.strip(): # Only process non-empty text | |
| st.text(f"{i+1}. {text}") | |
| new_text = st.text_input(f"Enter new text for '{text}':", value=text, key=f"new_text_{i}") | |
| if st.button(f"Replace '{text}'", key=f"replace_{i}"): | |
| edited_image = replace_text_in_image(edited_image, text, new_text) | |
| st.image(edited_image, caption='Edited Image', use_column_width=True) | |
| # Provide download option for the edited image | |
| buf = io.BytesIO() | |
| edited_image.save(buf, format="PNG") | |
| byte_im = buf.getvalue() | |
| st.download_button( | |
| label="Download edited image", | |
| data=byte_im, | |
| file_name="edited_image.png", | |
| mime="image/png" | |
| ) | |
| def poster_generator(): | |
| st.header("Generate Poster") | |
| poster_type = st.selectbox("Poster Type", ["Fashion", "Movie", "Event", "Advertisement", "Other"]) | |
| prompt = st.text_area("Prompt") | |
| num_images = st.number_input("Number of Images", min_value=1, max_value=5, value=1) | |
| quality_factor = st.number_input("Quality Factor", min_value=1, max_value=4, value=1) | |
| if st.button("Generate Images"): | |
| if poster_type == "Other": | |
| full_prompt = f"A colorful poster with the following elements: {prompt}" | |
| else: | |
| full_prompt = f"A colorful {poster_type.lower()} poster with the following elements: {prompt}" | |
| generated_images = [] | |
| for i in range(num_images): | |
| with st.spinner(f"Generating image {i+1}..."): | |
| logger.info(f"Generating image {i+1} with prompt: {full_prompt}") | |
| response = query({"inputs": full_prompt}) | |
| if isinstance(response, bytes): | |
| image = Image.open(io.BytesIO(response)) | |
| if quality_factor > 1: | |
| image = increase_image_quality(image, quality_factor) | |
| generated_images.append(image) | |
| else: | |
| st.error("Failed to generate image") | |
| # Display generated images | |
| for i, img in enumerate(generated_images): | |
| st.image(img, caption=f"Generated Poster {i+1}", use_column_width=True) | |
| # Provide download option for the generated image | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| byte_im = buf.getvalue() | |
| st.download_button( | |
| label=f"Download generated poster {i+1}", | |
| data=byte_im, | |
| file_name=f"generated_poster_{i+1}.png", | |
| mime="image/png" | |
| ) | |
| def main(): | |
| st.title("Image Processing App") | |
| app_mode = st.sidebar.selectbox("Choose the app mode", | |
| ["Image Text Replacer", "Poster Generator"]) | |
| if app_mode == "Image Text Replacer": | |
| image_text_replacer() | |
| elif app_mode == "Poster Generator": | |
| poster_generator() | |
| if __name__ == "__main__": | |
| main() |