Update app.py

app.py

@@ -5,6 +5,7 @@ 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')
@@ -13,16 +14,17 @@ model.load_weights('weights/RealESRGAN_x4.pth', download=True)
 
 def background_remover_and_upscale(input_image):
     # Remove background
-    output_img = remove(input_image)
+    output_img = remove(np.array(input_image))
     
     # Convert back to PIL Image for upscaling
-    output_img_pil = Image.fromarray(output_img)
+    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_pil)
+    sr_image = model.predict([output_img])
     
     # Return the upscaled image directly
-    return sr_image
+    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.
@@ -31,12 +33,16 @@ Simply upload your image and let the models do their work. Download the result i
 
 iface = gr.Interface(
     fn=background_remover_and_upscale,
-    inputs=gr.Image(type='pil', label="Upload Image"),
-    outputs="image",
+    inputs=gr.Image(type='pil'),
+    outputs=[gr.File("image")],
     examples=["woman.jpg", "groot.jpg"],
-    title="🖼 Background Remover and ✨ Image Upscaler",
+    title="Background Remover and Image Upscaler",
     description=description,
+    article="""This demo was created by KVI Kontent using Gradio.""",
+    allow_flagging=False,
+    live=True,
+    layout="horizontal",
     theme="soft"
 )
 
-iface.launch(share=True)
+iface.launch()