Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from openai import OpenAI | |
| import requests | |
| from PIL import Image | |
| import io | |
| import os | |
| from datetime import datetime | |
| def convert_to_png(image_file): | |
| """Convert uploaded image to PNG format and ensure it's under 4MB""" | |
| # Open the image using PIL | |
| image = Image.open(image_file) | |
| # Convert to RGBA if not already | |
| if image.mode in ('RGBA', 'RGB'): | |
| image = image.convert('RGBA') | |
| else: | |
| image = image.convert('RGB').convert('RGBA') | |
| # Save to bytes with PNG format | |
| byte_arr = io.BytesIO() | |
| image.save(byte_arr, format='PNG', optimize=True) | |
| byte_arr.seek(0) | |
| # Check if size is under 4MB | |
| if byte_arr.getbuffer().nbytes > 4 * 1024 * 1024: | |
| # If image is too large, resize it while maintaining aspect ratio | |
| while byte_arr.getbuffer().nbytes > 4 * 1024 * 1024: | |
| width, height = image.size | |
| new_width = int(width * 0.9) # Reduce by 10% | |
| new_height = int(height * 0.9) | |
| image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) | |
| byte_arr = io.BytesIO() | |
| image.save(byte_arr, format='PNG', optimize=True) | |
| byte_arr.seek(0) | |
| return byte_arr | |
| def validate_image(uploaded_file): | |
| """Validate image size and format""" | |
| if uploaded_file.size > 4 * 1024 * 1024: # 4MB in bytes | |
| return False, "File size must be less than 4MB" | |
| return True, "OK" | |
| def save_uploaded_file(uploaded_file, folder="uploads"): | |
| """Save uploaded file to a temporary folder and return the path""" | |
| if not os.path.exists(folder): | |
| os.makedirs(folder) | |
| # Generate a timestamp for unique filename | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| # Convert to PNG and validate size | |
| image_bytes = convert_to_png(uploaded_file) | |
| # Save the converted PNG | |
| file_path = os.path.join(folder, f"{timestamp}.png") | |
| with open(file_path, "wb") as f: | |
| f.write(image_bytes.getvalue()) | |
| return file_path | |
| def download_image(url, folder="generated"): | |
| """Download image from URL and save it""" | |
| if not os.path.exists(folder): | |
| os.makedirs(folder) | |
| response = requests.get(url) | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| file_path = os.path.join(folder, f"generated_image_{timestamp}.png") | |
| if response.status_code == 200: | |
| with open(file_path, "wb") as f: | |
| f.write(response.content) | |
| return file_path | |
| return None | |
| def main(): | |
| st.title("🎨 Image Editor with DALL-E 2") | |
| # Sidebar for API key | |
| with st.sidebar: | |
| st.header("Configuration") | |
| api_key = st.text_input("Enter your OpenAI API key", type="password") | |
| st.markdown(""" | |
| ### How to get an API key | |
| 1. Go to [OpenAI API Keys](https://platform.openai.com/api-keys) | |
| 2. Create a new secret key | |
| 3. Copy and paste it here | |
| """) | |
| # Size selection | |
| size_option = st.selectbox( | |
| "Select image size:", | |
| ["1024x1024", "512x512", "256x256"] | |
| ) | |
| # Number of images | |
| num_images = st.slider("Number of images to generate", 1, 4, 1) | |
| # Main content | |
| st.markdown(""" | |
| ### Requirements: | |
| - Original image and mask must be less than 4MB | |
| - Images will be automatically converted to PNG format | |
| - Images larger than 4MB will be automatically resized | |
| """) | |
| # File uploaders | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.subheader("Original Image") | |
| original_image = st.file_uploader( | |
| "Upload original image", | |
| type=["png", "jpg", "jpeg"] | |
| ) | |
| if original_image: | |
| # Validate image | |
| valid, message = validate_image(original_image) | |
| if not valid: | |
| st.warning(f"Original image: {message}. Image will be automatically resized.") | |
| try: | |
| # Display image preview | |
| image = Image.open(original_image) | |
| st.image(image, caption="Original Image", use_container_width=True) | |
| st.caption(f"Original size: {original_image.size/1024/1024:.2f}MB") | |
| except Exception as e: | |
| st.error(f"Error loading image: {str(e)}") | |
| with col2: | |
| st.subheader("Mask Image") | |
| mask_image = st.file_uploader( | |
| "Upload mask image", | |
| type=["png", "jpg", "jpeg"] | |
| ) | |
| if mask_image: | |
| # Validate image | |
| valid, message = validate_image(mask_image) | |
| if not valid: | |
| st.warning(f"Mask image: {message}. Image will be automatically resized.") | |
| try: | |
| # Display image preview | |
| image = Image.open(mask_image) | |
| st.image(image, caption="Mask Image", use_column_width=True) | |
| st.caption(f"Original size: {mask_image.size/1024/1024:.2f}MB") | |
| except Exception as e: | |
| st.error(f"Error loading mask: {str(e)}") | |
| # Prompt input | |
| prompt = st.text_area( | |
| "Enter your prompt:", | |
| placeholder="Describe the changes you want to make to the image...", | |
| help="Be specific about what you want to add or modify in the masked area" | |
| ) | |
| # Generate button | |
| if st.button("Generate Edited Image"): | |
| if not api_key: | |
| st.error("Please enter your OpenAI API key in the sidebar.") | |
| return | |
| if not original_image or not mask_image: | |
| st.error("Please upload both an original image and a mask image.") | |
| return | |
| if not prompt: | |
| st.error("Please enter a prompt describing the desired changes.") | |
| return | |
| try: | |
| with st.spinner("Processing images and generating edited version..."): | |
| # Save and convert uploaded files | |
| original_path = save_uploaded_file(original_image) | |
| mask_path = save_uploaded_file(mask_image) | |
| # Initialize OpenAI client | |
| client = OpenAI(api_key=api_key) | |
| # Make the API call | |
| response = client.images.edit( | |
| model="dall-e-2", | |
| image=open(original_path, "rb"), | |
| mask=open(mask_path, "rb"), | |
| prompt=prompt, | |
| n=num_images, | |
| size=size_option | |
| ) | |
| # Display results | |
| st.subheader("Generated Images") | |
| cols = st.columns(num_images) | |
| for idx, image_data in enumerate(response.data): | |
| # Download and save the generated image | |
| saved_image_path = download_image(image_data.url) | |
| if saved_image_path: | |
| with cols[idx]: | |
| st.image(saved_image_path, caption=f"Generated Image {idx+1}") | |
| # Add download button for each image | |
| with open(saved_image_path, "rb") as file: | |
| st.download_button( | |
| label=f"Download Image {idx+1}", | |
| data=file, | |
| file_name=f"edited_image_{idx+1}.png", | |
| mime="image/png" | |
| ) | |
| # Cleanup temporary files | |
| for path in [original_path, mask_path]: | |
| if os.path.exists(path): | |
| os.remove(path) | |
| except Exception as e: | |
| st.error(f"An error occurred: {str(e)}") | |
| finally: | |
| # Cleanup temporary files | |
| for path in [original_path, mask_path]: | |
| if 'path' in locals() and os.path.exists(path): | |
| try: | |
| os.remove(path) | |
| except Exception: | |
| pass | |
| if __name__ == "__main__": | |
| main() |