Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| import base64 | |
| from PIL import Image | |
| import gradio as gr | |
| from openai import OpenAI | |
| # Initialize the OpenAI client | |
| OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") | |
| client = OpenAI(api_key=OPENAI_API_KEY) | |
| def inference(input_image: Image.Image) -> Image.Image: | |
| """ | |
| Enhances a handwritten sketch into a clean, structured floor plan using OpenAI's image model. | |
| """ | |
| try: | |
| # Convert the PIL image to a PNG byte buffer | |
| img_byte_arr = io.BytesIO() | |
| input_image.save(img_byte_arr, format="PNG") | |
| img_byte_arr.seek(0) | |
| img_byte_arr.name = "input.png" # ✅ Give it a valid filename | |
| prompt = """ | |
| Enhance this handwritten sketch into a clean, structured, computer-drawn floor plan design. | |
| """ | |
| # Use the in-memory file directly | |
| result = client.images.edit( | |
| model="gpt-image-1", | |
| image=[img_byte_arr], | |
| prompt=prompt | |
| ) | |
| # Decode the returned base64 image | |
| image_base64 = result.data[0].b64_json | |
| image_bytes = base64.b64decode(image_base64) | |
| # Convert bytes back to PIL Image | |
| enhanced_image = Image.open(io.BytesIO(image_bytes)) | |
| return enhanced_image | |
| except Exception as e: | |
| print(f"Error during image enhancement: {e}") | |
| return input_image # fallback to original image | |
| # Gradio interface | |
| iface = gr.Interface( | |
| fn=inference, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Image(type="pil"), | |
| title="Sketch to Floor Plan Pix2Pix", | |
| description="Upload a sketch image and the model will enhance it into a structured floor plan." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |