Duplicate from szk1ck/image-matting

Co-authored-by: Akihiro SUZUKI <szk1ck@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Image Matting
+emoji: 🐢
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: 3.20.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: szk1ck/image-matting
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,263 @@
+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() 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"
+    )

requirements.txt

@@ -0,0 +1,94 @@
+aiofiles==23.1.0
+aiohttp==3.8.4
+aiosignal==1.3.1
+altair==4.2.2
+anyio==3.6.2
+appdirs==1.4.4
+async-timeout==4.0.2
+asyncer==0.0.2
+attrs==22.2.0
+blessed==1.20.0
+certifi==2022.12.7
+charset-normalizer==3.1.0
+click==8.1.3
+codefind==0.1.3
+coloredlogs==15.0.1
+contourpy==1.0.7
+cycler==0.11.0
+entrypoints==0.4
+fastapi==0.87.0
+ffmpy==0.3.0
+filetype==1.2.0
+flatbuffers==23.3.3
+fonttools==4.39.0
+frozenlist==1.3.3
+fsspec==2023.3.0
+gradio==3.20.1
+h11==0.14.0
+httpcore==0.16.3
+httpx==0.23.3
+humanfriendly==10.0
+idna==3.4
+ImageHash==4.3.1
+imageio==2.26.0
+importlib-resources==5.12.0
+Jinja2==3.1.2
+jsonschema==4.17.3
+jurigged==0.5.5
+kiwisolver==1.4.4
+lazy_loader==0.1
+linkify-it-py==2.0.0
+llvmlite==0.39.1
+markdown-it-py==2.2.0
+MarkupSafe==2.1.2
+matplotlib==3.7.1
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.0.4
+networkx==3.0
+numba==0.56.4
+numpy==1.23.5
+onnxruntime==1.13.1
+opencv-python-headless==4.6.0.66
+orjson==3.8.7
+ovld==0.3.2
+packaging==23.0
+pandas==1.5.3
+Pillow==9.3.0
+pooch==1.6.0
+protobuf==4.22.1
+pycryptodome==3.17
+pydantic==1.10.6
+pydub==0.25.1
+PyMatting==1.1.8
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytz==2022.7.1
+PyWavelets==1.4.1
+PyYAML==6.0
+rembg==2.0.30
+requests==2.28.2
+rfc3986==1.5.0
+scikit-image==0.19.3
+scipy==1.9.3
+six==1.16.0
+sniffio==1.3.0
+starlette==0.21.0
+sympy==1.11.1
+tifffile==2023.2.28
+toolz==0.12.0
+torch==1.13.1
+torchvision==0.14.1
+tqdm==4.64.1
+typing_extensions==4.5.0
+uc-micro-py==1.0.1
+urllib3==1.26.15
+uvicorn==0.20.0
+watchdog==2.1.9
+wcwidth==0.2.6
+websockets==10.4
+yarl==1.8.2
+zipp==3.15.0

utils/functions.py

@@ -0,0 +1,70 @@
+import os, random
+from logging import getLogger, StreamHandler, DEBUG
+logger = getLogger(__name__)
+handler = StreamHandler(); handler.setLevel(DEBUG)
+logger.setLevel(DEBUG)
+logger.addHandler(handler)
+logger.propagate = False
+
+
+def get_random_name():
+    famous_painters = [
+    "Leonardo", "DaVinci",
+    "Michelangelo",
+    "Pablo", "Picasso",
+    "Vincent", "VanGogh",
+    "Rembrandt", "VanRijn",
+    "Claude", "Monet",
+    "Salvador", "Dali",
+    "Jackson", "Pollock",
+    "Andy", "Warhol",
+    "Henri", "Matisse",
+    "Georgia", "Keeffe",
+    "Edvard", "Munch",
+    "Wassily", "Kandinsky",
+    "Gustav", "Klimt",
+    "Rene", "Magritte",
+    "Frida", "Kahlo",
+    "Edgar", "Degas",
+    "Johannes", "Vermeer",
+    "Paul", "Cezanne",
+    "Marc", "Chagall",
+    ]
+
+    random_painter = random.choice(famous_painters)
+
+    # 4桁の乱数を生成
+    rand_num = random.randint(1000, 9999)
+
+    return random_painter+str(rand_num)
+
+
+def complete(work_dir):
+    work_dir = work_dir
+    # logger.debug(f"complete :", work_dir)
+    return work_dir
+
+
+def clean(text_output):
+    # logger.debug(f"text_output : {text_output}")
+
+    if text_output!="idle_state":
+        logger.info(f"clean up : {text_output}.zip")
+        os.remove(f"{text_output}.zip")
+        return "idle_state"
+    else:
+        logger.info(f"reset")
+        return "idle_state"
+
+
+
+def clean_by_name(text_output):
+    # logger.debug(f"text_output : {text_output}")
+    if text_output!="idle_state":
+        text_output, dir_name = text_output.split("+")
+        logger.info(f"clean up : {dir_name}.zip")
+        os.remove(f"{dir_name}.zip")
+        return "idle_state"
+    else:
+        logger.info(f"reset")
+        return "idle_state"