Add files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* 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
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.pre-commit-config.yaml

@@ -0,0 +1,50 @@
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.4.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ["--fix=lf"]
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.7.5
+  hooks:
+  - id: docformatter
+    args: ["--in-place"]
+- repo: https://github.com/pycqa/isort
+  rev: 5.12.0
+  hooks:
+    - id: isort
+      args: ["--profile", "black"]
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v1.5.1
+  hooks:
+    - id: mypy
+      args: ["--ignore-missing-imports"]
+      additional_dependencies: ["types-python-slugify", "types-requests", "types-PyYAML"]
+- repo: https://github.com/psf/black
+  rev: 23.9.1
+  hooks:
+    - id: black
+      language_version: python3.10
+      args: ["--line-length", "119"]
+- repo: https://github.com/kynan/nbstripout
+  rev: 0.6.1
+  hooks:
+    - id: nbstripout
+      args: ["--extra-keys", "metadata.interpreter metadata.kernelspec cell.metadata.pycharm"]
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.7.0
+  hooks:
+    - id: nbqa-black
+    - id: nbqa-pyupgrade
+      args: ["--py37-plus"]
+    - id: nbqa-isort
+      args: ["--float-to-top"]

.vscode/settings.json

@@ -0,0 +1,21 @@
+{
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter",
+        "editor.formatOnType": true,
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": true
+        }
+    },
+    "black-formatter.args": [
+        "--line-length=119"
+    ],
+    "isort.args": ["--profile", "black"],
+    "flake8.args": [
+        "--max-line-length=119"
+    ],
+    "ruff.args": [
+        "--line-length=119"
+    ],
+    "editor.formatOnSave": true,
+    "files.insertFinalNewline": true
+}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 hysts
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -5,8 +5,16 @@ colorFrom: indigo
 colorTo: blue
 sdk: gradio
 sdk_version: 3.44.4
+python_version: 3.10.12
 app_file: app.py
 pinned: false
+license: mit
+suggested_hardware: t4-small
+models:
+  - hustvl/vitmatte-small-distinctions-646
+  - hustvl/vitmatte-base-distinctions-646
+  - hustvl/vitmatte-small-composition-1k
+  - hustvl/vitmatte-base-composition-1k
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+
+import os
+
+import gradio as gr
+import numpy as np
+import PIL.Image
+import spaces
+import torch
+from transformers import VitMatteForImageMatting, VitMatteImageProcessor
+
+DESCRIPTION = "# [ViTMatte](https://github.com/hustvl/ViTMatte)"
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1500"))
+MODEL_ID = os.getenv("MODEL_ID", "hustvl/vitmatte-small-distinctions-646")
+
+processor = VitMatteImageProcessor.from_pretrained(MODEL_ID)
+model = VitMatteForImageMatting.from_pretrained(MODEL_ID).to(device)
+
+
+def check_image_size(image: PIL.Image.Image) -> None:
+    if max(image.size) > MAX_IMAGE_SIZE:
+        raise gr.Error(f"Image size is too large. Max image size is {MAX_IMAGE_SIZE} pixels.")
+
+
+def binarize_mask(mask: np.ndarray) -> np.ndarray:
+    mask[mask < 128] = 0
+    mask[mask > 0] = 1
+    return mask
+
+
+def update_trimap(foreground_mask: dict[str, np.ndarray], unknown_mask: dict[str, np.ndarray]) -> np.ndarray:
+    foreground = foreground_mask["mask"]
+    foreground = binarize_mask(foreground)
+
+    unknown = unknown_mask["mask"]
+    unknown = binarize_mask(unknown)
+
+    trimap = np.zeros_like(foreground)
+    trimap[unknown > 0] = 128
+    trimap[foreground > 0] = 255
+    return trimap
+
+
+@spaces.GPU
+@torch.inference_mode()
+def run(image: PIL.Image.Image, trimap: PIL.Image.Image) -> tuple[PIL.Image.Image, PIL.Image.Image]:
+    if image.size != trimap.size:
+        raise gr.Error("Image and trimap must have the same size.")
+    if max(image.size) > MAX_IMAGE_SIZE:
+        raise gr.Error(f"Image size is too large. Max image size is {MAX_IMAGE_SIZE} pixels.")
+    if image.mode != "RGB":
+        raise gr.Error("Image must be RGB.")
+    if trimap.mode != "L":
+        raise gr.Error("Trimap must be grayscale.")
+
+    pixel_values = processor(images=image, trimaps=trimap, return_tensors="pt").to(device).pixel_values
+    out = model(pixel_values=pixel_values)
+    alpha = out.alphas[0, 0].to("cpu").numpy()
+
+    w, h = image.size
+    alpha = alpha[:h, :w]
+
+    foreground = np.array(image).astype(float) / 255 * alpha[:, :, None] + (1 - alpha[:, :, None])
+    foreground = (foreground * 255).astype(np.uint8)
+    foreground = PIL.Image.fromarray(foreground)
+
+    return alpha, foreground
+
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Box():
+                image = gr.Image(label="Input image", type="pil", height=500)
+                with gr.Tabs():
+                    with gr.Tab(label="Trimap"):
+                        trimap = gr.Image(label="Trimap", type="pil", image_mode="L", height=500)
+                    with gr.Tab(label="Draw trimap"):
+                        load_image_button = gr.Button("Load image")
+                        foreground_mask = gr.Image(
+                            label="Foreground",
+                            tool="sketch",
+                            type="numpy",
+                            brush_color="green",
+                            mask_opacity=0.7,
+                            height=500,
+                        )
+                        unknown_mask = gr.Image(
+                            label="Unkown",
+                            tool="sketch",
+                            type="numpy",
+                            brush_color="green",
+                            mask_opacity=0.7,
+                            height=500,
+                        )
+                        set_trimap_button = gr.Button("Set trimap")
+                run_button = gr.Button("Run")
+        with gr.Column():
+            with gr.Box():
+                out_alpha = gr.Image(label="Alpha", height=500)
+                out_foreground = gr.Image(label="Foreground", height=500)
+
+    gr.Examples(
+        examples=[
+            ["assets/bulb_rgb.png", "assets/bulb_trimap.png"],
+            ["assets/retriever_rgb.png", "assets/retriever_trimap.png"],
+        ],
+        inputs=[image, trimap],
+        outputs=[out_alpha, out_foreground],
+        fn=run,
+        cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
+    )
+
+    image.change(
+        fn=check_image_size,
+        inputs=image,
+        queue=False,
+        api_name=False,
+    )
+    load_image_button.click(
+        fn=lambda image: (image, image),
+        inputs=image,
+        outputs=[foreground_mask, unknown_mask],
+        queue=False,
+        api_name=False,
+    )
+    set_trimap_button.click(
+        fn=update_trimap,
+        inputs=[foreground_mask, unknown_mask],
+        outputs=trimap,
+        queue=False,
+        api_name=False,
+    )
+    run_button.click(
+        fn=run,
+        inputs=[image, trimap],
+        outputs=[out_alpha, out_foreground],
+        api_name="run",
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()

assets/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Hust Vision Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

assets/README.md

@@ -0,0 +1,4 @@
+# About these images
+
+These images were obtained from [this link](https://github.com/hustvl/ViTMatte/tree/f48208ceb0d172ab0c1c528bba33cc1159596407/demo).
+The original repository is licensed under the MIT license.

requirements.txt

@@ -0,0 +1,4 @@
+git+https://github.com/huggingface/transformers@8e3980a
+Pillow==10.0.1
+torch==2.0.0
+torchvision==0.15.1

style.css

@@ -0,0 +1,10 @@
+h1 {
+  text-align: center;
+}
+
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+}