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
from utils import functions
from utils.functions import complete, clean, clean_by_name, get_random_name

import os, sys
import argparse

from logging import getLogger, StreamHandler, DEBUG
logger = getLogger(__name__)
handler = StreamHandler(); handler.setLevel(DEBUG)
logger.setLevel(DEBUG)
logger.addHandler(handler)
logger.propagate = False


        
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 from_zip(inputs):
    work_dir = get_random_name()
    os.makedirs(work_dir, 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)
                        # logger.debug(f"image name : {image_name}")
                        category_dir = image_name.split("/")[0]
                        # image_name = image_name.split("/")[1]
                        os.makedirs(f"{work_dir}/{category_dir}", exist_ok=True)
                        image_data_dict[image_name] = image_array

                except Exception as e:
                    logger.info(f"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"{work_dir}/{image_name}")
    
    shutil.make_archive(work_dir, "zip", work_dir)
    shutil.rmtree(work_dir)
    
    return f"{work_dir}.zip", complete(work_dir)


def from_image_files(images, text_class_name):

    if not text_class_name=="":
        dir_name = text_class_name
    else:
        dir_name = functions.get_random_name()

    os.makedirs(dir_name, exist_ok=True)

    for image in images:
        image_name = image.name
        # logger.debug(f"image name : {image_name}")
    
        # 読み込み
        image_data = np.array(Image.open(image_name))

        output = remove(image_data)        
        output_pil = Image.fromarray(output)
        # Remove margins 
        cropped_image = output_pil.crop(output_pil.getbbox())

        image_name = image_name.split("/")[-1]
        image_name = image_name[:image_name.find("_", image_name.find("_") + 1)] + ".png"
        # logger.debug(f"save image name : {image_name}")
        cropped_image.save(f"{dir_name}/{image_name}")
    
    shutil.make_archive(f"{dir_name}", "zip", f"{dir_name}")
    shutil.rmtree(f"{dir_name}")
    
    return f"{dir_name}.zip", complete("complete")+"+"+dir_name



if __name__=="__main__":

       
    with gr.Blocks(
        css=".gradio-container {background: url('file=./assets/codioful-formerly-gradienta-bKESVqfxass-unsplash.jpg')}"
    ) as demo:

        
        with gr.Tab("Images"):
            gr.Markdown(
            """
            <center><h1>Image Matting using U<sup>2</sup>-Net</h1></center>
            """
            )        
            with gr.Row():
                gr.Markdown(
                    """
                    ### Input Image Files
                    <img src='file/assets/input_images.png' width="90%" height="90%">
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    """
                )        
            
            with gr.Row():
                with gr.Column():
                    text_class_name = gr.Textbox(label="Class Name", value="", placeholder="cat")
                    image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False)
            
            btn = gr.Button("Run!")
            
            btn.click(
                fn=from_image_files,
                inputs=[image_input, text_class_name],
                outputs=[image_output, text_output]
                )
            text_output.change(
                fn=clean_by_name,
                inputs=text_output,
                outputs=text_output

            )


        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%">
                    Zip file can include multiple directories.
                    """
                )        
                gr.Markdown(
                    """
                    ### Output Zip File
                    <img src='file/assets/output_zip.png' width="85%" height="85%">
                    If input has multiple directories, output has the same multiple diretocories.
                    """
                )        
            
            with gr.Row():
                image_input = gr.File(file_count="multiple")
                image_output = gr.File()
                text_output = gr.Textbox(visible=False, value="idle_state")
            
            btn = gr.Button("Run!")
            
            btn.click(
                fn=from_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"
    )