app.py

import os
import uuid
import base64
import requests
from PIL import Image
from io import BytesIO
from pathlib import Path
import gradio as gr
from gradio_imageslider import ImageSlider  # Ensure this library is installed
from dotenv import load_dotenv

import config

# Load environment variables from the .env file
load_dotenv()

# Get API key from environment variable
api_key = os.getenv('API_KEY')

# Funzione per chiamare l'endpoint di predizione FastAPI
def process_image(input_image_editor):
    input_image = input_image_editor['background']
    mask_image = input_image_editor['layers'][0]

    # Converti le immagini in base64
    buffered_input = BytesIO()
    input_image.save(buffered_input, format="PNG")
    input_image_base64 = base64.b64encode(buffered_input.getvalue()).decode()

    buffered_mask = BytesIO()
    mask_image.save(buffered_mask, format="PNG")
    mask_image_base64 = base64.b64encode(buffered_mask.getvalue()).decode()

    # Prepara il payload per la richiesta POST
    payload = {
        "input_image_editor": {
            "background": input_image_base64,
            "layers": [mask_image_base64]
        }
    }

    # Effettua la richiesta POST al backend FastAPI
    response = requests.post(
        os.getenv('BACKEND_URL') + "/predict/",
        headers={"access_token": api_key},
        json=payload
    )

    if response.status_code == 200:
        result = response.json()
        processed_image_base64 = result['processed_image']
        processed_image = Image.open(BytesIO(base64.b64decode(processed_image_base64)))

        # Save the processed image
        output_folder = Path("output")  # Make sure this folder exists or create it
        output_folder.mkdir(parents=True, exist_ok=True)
        image_path = output_folder / f"no_bg_image_{uuid.uuid4().hex}.png"
        processed_image.save(image_path)

        return (processed_image, input_image), str(image_path)

    else:
        raise Exception(f"Request failed with status code {response.status_code}")



# Define inputs and outputs for the Gradio interface
image = gr.ImageEditor(
                label='Image',
                type='pil',
                sources=["upload", "webcam"],
                image_mode='RGB',
                layers=False,
                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed")
            )
output_slider = ImageSlider(label="Processed photo", type="pil")

demo = gr.Interface(
    fn=process_image,
    inputs=image,
    outputs=[output_slider, gr.File(label="output png file")],
    title=config.TITLE,
    description=config.DESCRIPTION,
    article=config.BUY_ME_A_COFFE
)

#Center the title and description using custom CSS
demo.css = """  
    .interface-title {  
        text-align: center;  
    }  
    .interface-description {  
        text-align: center;  
    }  
"""

demo.launch(debug=False, show_error=True, share=True)