import gradio as gr
from rembg import remove
from PIL import Image
import io
import torch
import numpy as np
from RealESRGAN import RealESRGAN
import cv2

# Initialize RealESRGAN model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RealESRGAN(device, scale=4)
model.load_weights('weights/RealESRGAN_x4.pth', download=True)

def background_remover_and_upscale(input_image):
    # Remove background
    output_img = remove(np.array(input_image))
    
    # Convert back to PIL Image for upscaling
    output_img_rgb = cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB)
    output_img_pil = Image.fromarray(output_img_rgb)
    
    # Upscale image
    sr_image = model.predict([output_img])
    
    # Return the upscaled image directly
    return sr_image[0]

description = """
✨ This app uses the rembg library to remove the background from uploaded images and then the RealESRGAN model to upscale the image.
Simply upload your image and let the models do their work. Download the result immediately after!
"""

iface = gr.Interface(
    fn=background_remover_and_upscale,
    inputs=gr.Image(type='pil'),
    outputs="image",
    examples=["woman.jpg", "groot.jpg"],
    title="Background Remover and Image Upscaler",
    description=description,
    article="""This demo was created by KVI Kontent using Gradio.""",
    theme="soft"
)

iface.launch()