Update app.py

app.py

@@ -1,55 +1,20 @@
-import os
-import sys
-import logging
+import gradio as gr
 from rembg import remove
-from PIL import Image
 
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(message)s')
-
-def process_folder(input_folder, output_folder):
-    # Ensure the output folder exists
-    os.makedirs(output_folder, exist_ok=True)
-
-    # Define supported image extensions
-    supported_extensions = ('.png', '.jpg', '.jpeg')
-
-    # Iterate over all files in the input folder
-    for filename in os.listdir(input_folder):
-        # Check if the file is an image
-        if filename.lower().endswith(supported_extensions):
-            try:
-                # Load the image
-                image_path = os.path.join(input_folder, filename)
-                image = Image.open(image_path)
-
-                # Remove the background
-                image_no_bg = remove(image)
-
-                # Save the output image with the prefix 'bg-rm_'
-                output_filename = 'bg-rm_' + filename
-                output_path = os.path.join(output_folder, output_filename)
-                image_no_bg.save(output_path)
-
-                # Log success
-                logging.info(f'Processed: {filename}')
-
-            except Exception as e:
-                # Log error
-                logging.error(f'Error processing {filename}: {e}')
-
-def main():
-    # Check command-line arguments
-    if len(sys.argv) != 3:
-        print(f'Usage: {sys.argv[0]} input_folder output_folder')
-        sys.exit(1)
-
-    # Get input and output folders from command-line arguments
-    input_folder = sys.argv[1]
-    output_folder = sys.argv[2]
-
-    # Process the images in the input folder
-    process_folder(input_folder, output_folder)
-
-if __name__ == '__main__':
-    main()
+title = "Vanish Background"
+description = "Remove background for any image. To use it, simply upload your image(s) and wait. Read more at the link below of official documentation."
+article = "<p style='text-align: center;'><a href='https://github.com/danielgatis/rembg' target='_blank'>Github Repo</a></p>"
+
+def segment(images):
+    # Process each image in the list and remove the background
+    return [remove(image) for image in images]
+
+demo = gr.Interface(
+    fn=segment,
+    inputs=gr.inputs.Image(type="file", label="Upload Image(s)", multiple=True),
+    outputs=gr.outputs.Image(type="file", label="Image(s) with Background Removed"),
+    title=title,
+    description=description,
+    article=article
+)
+demo.launch()