app.py

from zipfile import ZipFile
import numpy as np

import os, shutil

import gradio as gr
from PIL import Image
import sys, os
from rembg import remove


def complete():
    return "complete"


def clean(text_output):
    if text_output=="complete":
        print("clean up")
        os.remove("objects.zip")
        return ""
    else:
        pass
    
    

def run_rembg(img):
    output = remove(img)
    output_pil = Image.fromarray(output)

    # Remove margins 
    cropped_image = output_pil.crop(output_pil.getbbox())

    return cropped_image

def load_zip(inputs):
    os.makedirs("objects", exist_ok=True)

    image_data_dict = {}
    with ZipFile(inputs[0].name, "r") as zip_file:
        image_names = zip_file.namelist()

        prefix = ""
        for name in image_names:
            if prefix=="":
                prefix = name.split("/")[0]
            else:
                break

        image_files = []      
          
        for image_name in image_names:
            if image_name[-3:] in "pngjpg":
                try:
                    with zip_file.open(image_name) as f:
                        image = Image.open(f)
                        image_files.append(image_name)
                        
                        image_array = np.array(image)
                        image_name = image_name.split("/")[1]
                        image_data_dict[image_name] = image_array

                except Exception as e:
                    print("Exception : ", e)
    
    
    for image_name, image_data in image_data_dict.items():
        output = remove(image_data)        
        output_pil = Image.fromarray(output)
        # Remove margins 
        cropped_image = output_pil.crop(output_pil.getbbox())

        image_name = image_name.replace("jpg", "png")
        cropped_image.save(f"objects/{image_name}")
    
    shutil.make_archive("objects", "zip", "objects")
    shutil.rmtree("objects")
    
    return "objects.zip", complete()



if __name__=="__main__":

       
    with gr.Blocks() as demo:

        with gr.Tab("Zip"):
            gr.Markdown(
            """
            <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
            """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Zip File
                    <img src='file/assets/input_zip.png' width="85%" height="85%">
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    """
                )        
            
            with gr.Row():
                image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False)
            
            btn = gr.Button("Run!")
            
            btn.click(
                fn=load_zip,
                inputs=image_input,
                outputs=[image_output, text_output]
                )
            text_output.change(
                fn=clean,
                inputs=text_output,
                outputs=text_output

            )


        with gr.Tab("Image"):
            gr.Markdown(
                """
                <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
                """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Image
                    <img src='file/assets/cat-3038243_1280.jpg' width="50%" height="50%">
                    """
                )        
                
                gr.Markdown(
                    """
                    ### Output Image
                    <img src='file/assets/objects/cat-3038243_1280.png' width="20%" height="20%">
                    """
                )        
            with gr.Row():
                image_input = gr.Image(type="numpy")
                image_output = gr.Image(type="pil")
            
            btn = gr.Button("Run!")
        
        btn.click(
            fn=run_rembg,
            inputs=image_input,
            outputs=image_output,
            api_name="imageMatting"
            )

        gr.Markdown(
            """
            ---
            <b>Acknowledgments</b>
            - Library
              - Library Git hub : [danielgatis/rembg](https://github.com/danielgatis/rembg)
              - Cloned on 2023/3/12  
            - Algorithm        
              - Library Git hub : [U<sup>2</sup>-Net](https://github.com/xuebinqin/U-2-Net)  
            - Image
              - Cat Image from [Pixabay](https://pixabay.com/images/id-3038243/)
            """
        )

    demo.launch(
        favicon_path="./assets/ハサミのフリーアイコン.png"
    )