app.py

import gradio as gr
from PIL import Image
from transparent_background import Remover
import numpy as np

# Initialize the model globally
remover = Remover()

def process_image(input_image, torchscript_jit, output_type):
    global remover
    if torchscript_jit == "on" and not isinstance(remover, Remover(jit=True).__class__):
        remover = Remover(jit=True)
    elif torchscript_jit == "default" and not isinstance(remover, Remover().__class__):
        remover = Remover()
    
    if output_type == "Mask only":
        # Process the image and get only the mask
        output = remover.process(input_image, type='map')
        if isinstance(output, Image.Image):
            # If output is already a PIL Image, convert to grayscale
            mask = output.convert('L')
        else:
            # If output is a numpy array, convert to PIL Image
            mask = Image.fromarray((output * 255).astype(np.uint8), mode='L')
        return mask
    else:
        # Process the image and return the RGBA result
        output = remover.process(input_image, type='rgba')
        return output

iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Radio(["default", "on"], label="TorchScript JIT", value="default"),
        gr.Radio(["Default", "Mask only"], label="Output Type", value="Default")
    ],
    outputs=gr.Image(type="pil", label="Output Image"),
    title="Inspyrenet Background Remover",
    description="Remove the background from an image using Inspyrenet. Choose 'Mask only' for a black and white mask, or 'Default' for the image with transparent background.",
    theme='bethecloud/storj_theme',
)

if __name__ == "__main__":
    iface.launch()