File size: 6,165 Bytes
d8a9fb0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c59606b
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
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()