print(">>> Starting app.py")
import torch
import torchvision
from PIL import Image
import torchvision.transforms as T
import numpy as np
import gradio as gr

# Load pretrained DeepLabV3 model once
model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)
model.eval()

# Define background removal function
def remove_bg(img):
    # Preprocess image
    transform = T.Compose([
        T.Resize(520),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
    ])
    input_tensor = transform(img).unsqueeze(0)

    # Run inference
    with torch.no_grad():
        output = model(input_tensor)['out'][0]

    mask = output.argmax(0).byte().cpu().numpy()

    # Resize mask to original image
    img_np = np.array(img)
    mask_resized = Image.fromarray(mask).resize((img_np.shape[1], img_np.shape[0]))
    mask_np = np.array(mask_resized)

    # Apply mask: keep object, set background white
    removed_bg = img_np.copy()
    removed_bg[mask_np == 0] = [255, 255, 255]

    return Image.fromarray(removed_bg)

# Create Gradio interface
demo = gr.Interface(
    fn=remove_bg,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Background Remover",
    description="Upload an image and get the background removed instantly!"
)

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