HF integration +uv + cli

2026-03-13 00:32:46 +01:00
parent 400e00418f
commit 9839d72332
7 changed files with 3295 additions and 50 deletions
@@ -6,3 +6,5 @@ results/
 test_sample/
 pretrained_models/
 data/
+notebooks/
+.venv
@@ -49,6 +49,7 @@


 ## 📮 Update
+- [2026.03] Add uv and huggingface support for easy installation and usage.
 - [2026.03] Release inference codes, evaluation codes, and gradio demo.
 - [2025.12] This repo is created.

@@ -65,23 +66,35 @@
 ![overall_structure](assets/matanyone1vs2.jpg)

 ## 🔧 Installation
-1. Clone Repo
-    ```bash
-    git clone https://github.com/pq-yang/MatAnyone2
-    cd MatAnyone2
-    ```

-2. Create Conda Environment and Install Dependencies
-    ```bash
-    # create new conda env
-    conda create -n matanyone2 python=3.10 -y
-    conda activate matanyone2
+### From PyPI / Source
+```bash
+# install from repo
+pip install git+https://github.com/pq-yang/MatAnyone2.git#egg=matanyone2

-    # install python dependencies
-    pip install -e .
-    # [optional] install python dependencies for gradio demo
-    pip3 install -r hugging_face/requirements.txt
-    ```
+# or install optional extras
+pip install  git+https://github.com/pq-yang/MatAnyone2.git#egg=matanyone2[gui]    # Gradio demo + PySide6
+pip install  git+https://github.com/pq-yang/MatAnyone2.git#egg=matanyone2[dev]    # development / evaluation tools
+pip install  git+https://github.com/pq-yang/MatAnyone2.git#egg=matanyone2[all]    # everything
+```
+
+### Conda
+```bash
+conda create -n matanyone2 python=3.10 -y
+conda activate matanyone2
+pip install  git+https://github.com/pq-yang/MatAnyone2.git#egg=matanyone2[all]
+```
+
+### uv (recommended)
+```bash
+# create a new project and add matanyone2
+uv init my-matting-project && cd my-matting-project
+uv add matanyone2@git+https://github.com/pq-yang/MatAnyone2.git
+
+# or with optional extras
+uv add matanyone2[gui]@git+https://github.com/pq-yang/MatAnyone2.git
+uv add matanyone2[all]@git+https://github.com/pq-yang/MatAnyone2.git
+```

 ## 🔥 Inference

@@ -109,15 +122,33 @@ Run the following command to try it out:

 ```shell
 # intput format: video folder
-python inference_matanyone2.py -i inputs/video/test-sample1 -m inputs/mask/test-sample1.png
+matanyone2 -i inputs/video/test-sample1 -m inputs/mask/test-sample1.png

 # intput format: mp4
-python inference_matanyone2.py -i inputs/video/test-sample2.mp4 -m inputs/mask/test-sample2.png
+matanyone2 -i inputs/video/test-sample2.mp4 -m inputs/mask/test-sample2.png

+# or via python
+python inference_matanyone2.py -i inputs/video/test-sample1 -m inputs/mask/test-sample1.png
 ```
 The results will be saved in the `results` folder, including the foreground output video and the alpha output video.
- If you want to save the results as per-frame images, you can set `--save_image`.
- If you want to set a limit for the maximum input resolution, you can set `--max_size`, and the video will be downsampled if min(w, h) exceeds. By default, we don't set the limit.
+
+### Python API (recommended 🔥)
+You can load the model directly from Hugging Face using `from_pretrained` and run inference programmatically:
+
+```python
+from matanyone2 import MatAnyone2, InferenceCore
+
+model = MatAnyone2.from_pretrained("not-lain/matanyone2")
+processor = InferenceCore(model, device="cuda:0")
+processor.process_video(
+    input_path="inputs/video/test-sample2.mp4",
+    mask_path="inputs/mask/test-sample2.png",
+    output_path="results",
+)
+``` 
+- If you want to save the results as per-frame images, you can set `--save-image`.
+- If you want to set a limit for the maximum input resolution, you can set `--max-size`, and the video will be downsampled if min(w, h) exceeds. By default, we don't set the limit.
+- Run `matanyone2 --help` for a full list of options.

 ## 🎪 Interactive Demo
 To get rid of the preparation for first-frame segmentation mask, we prepare a gradio demo on [hugging face](https://huggingface.co/spaces/PeiqingYang/MatAnyone2) and could also **launch locally**. Just drop your video/image, assign the target masks with a few clicks, and get the the matting results!
@@ -127,7 +158,7 @@ To get rid of the preparation for first-frame segmentation mask, we prepare a gr
 ```shell
 cd hugging_face

-# install python dependencies
+# install GUI dependencies (if not already installed via pip install -e ".[gui]")
 pip3 install -r requirements.txt # FFmpeg required

 # launch the demo
@@ -0,0 +1,171 @@
+import os
+from typing import Annotated
+
+import cv2
+import tqdm
+import typer
+import imageio
+import numpy as np
+from PIL import Image
+
+import torch
+import torch.nn.functional as F
+
+from matanyone2.utils.download_util import load_file_from_url
+from matanyone2.utils.inference_utils import gen_dilate, gen_erosion, read_frame_from_videos
+from matanyone2.inference.inference_core import InferenceCore
+from matanyone2.utils.get_default_model import get_matanyone2_model
+from matanyone2.utils.device import get_default_device, safe_autocast_decorator
+
+import warnings
+warnings.filterwarnings("ignore")
+
+app = typer.Typer(help="MatAnyone2 video matting inference")
+
+
+@torch.inference_mode()
+@safe_autocast_decorator()
+def run_inference(input_path, mask_path, output_path, ckpt_path, n_warmup=10, r_erode=10,
+                  r_dilate=10, suffix="", save_image=False, max_size=-1):
+    device = get_default_device()
+
+    # download ckpt for the first inference
+    pretrain_model_url = "https://github.com/pq-yang/MatAnyone2/releases/download/v1.0.0/matanyone2.pth"
+    ckpt_path = load_file_from_url(pretrain_model_url, 'pretrained_models')
+
+    # load MatAnyone model
+    matanyone2 = get_matanyone2_model(ckpt_path, device)
+
+    # init inference processor
+    processor = InferenceCore(matanyone2, cfg=matanyone2.cfg)
+
+    # inference parameters
+    r_erode = int(r_erode)
+    r_dilate = int(r_dilate)
+    n_warmup = int(n_warmup)
+    max_size = int(max_size)
+
+    # load input frames
+    vframes, fps, length, video_name = read_frame_from_videos(input_path)
+    repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1)
+    vframes = torch.cat([repeated_frames, vframes], dim=0).float()
+    length += n_warmup
+
+    # resize if needed
+    if max_size > 0:
+        h, w = vframes.shape[-2:]
+        min_side = min(h, w)
+        if min_side > max_size:
+            new_h = int(h / min_side * max_size)
+            new_w = int(w / min_side * max_size)
+            vframes = F.interpolate(vframes, size=(new_h, new_w), mode="area")
+            print(f'Resize to {new_h}x{new_w} for processing...')
+
+    # set output paths
+    os.makedirs(output_path, exist_ok=True)
+    if suffix != "":
+        video_name = f'{video_name}_{suffix}'
+    if save_image:
+        os.makedirs(f'{output_path}/{video_name}', exist_ok=True)
+        os.makedirs(f'{output_path}/{video_name}/pha', exist_ok=True)
+        os.makedirs(f'{output_path}/{video_name}/fgr', exist_ok=True)
+
+    # load the first-frame mask
+    mask = Image.open(mask_path).convert('L')
+    mask = np.array(mask)
+
+    bgr = (np.array([120, 255, 155], dtype=np.float32) / 255).reshape((1, 1, 3))
+    objects = [1]
+
+    # [optional] erode & dilate
+    if r_dilate > 0:
+        mask = gen_dilate(mask, r_dilate, r_dilate)
+    if r_erode > 0:
+        mask = gen_erosion(mask, r_erode, r_erode)
+
+    mask = torch.from_numpy(mask).float().to(device)
+
+    if max_size > 0:
+        mask = F.interpolate(mask.unsqueeze(0).unsqueeze(0), size=(new_h, new_w), mode="nearest")
+        mask = mask[0, 0]
+
+    # inference start
+    phas = []
+    fgrs = []
+    for ti in tqdm.tqdm(range(length)):
+        image = vframes[ti]
+        image_np = np.array(image.permute(1, 2, 0))
+        image = (image / 255.).float().to(device)
+
+        if ti == 0:
+            output_prob = processor.step(image, mask, objects=objects)
+            output_prob = processor.step(image, first_frame_pred=True)
+        else:
+            if ti <= n_warmup:
+                output_prob = processor.step(image, first_frame_pred=True)
+            else:
+                output_prob = processor.step(image)
+
+        mask = processor.output_prob_to_mask(output_prob)
+
+        pha = mask.unsqueeze(2).cpu().numpy()
+        com_np = image_np / 255. * pha + bgr * (1 - pha)
+
+        if ti > (n_warmup - 1):
+            com_np = np.round(np.clip(com_np * 255.0, 0, 255)).astype(np.uint8)
+            pha = np.round(np.clip(pha * 255.0, 0, 255)).astype(np.uint8)
+            fgrs.append(com_np)
+            phas.append(pha)
+            if save_image:
+                cv2.imwrite(f'{output_path}/{video_name}/fgr/{str(ti - n_warmup).zfill(4)}.png',
+                            com_np[..., [2, 1, 0]])
+                cv2.imwrite(f'{output_path}/{video_name}/pha/{str(ti - n_warmup).zfill(4)}.png',
+                            pha)
+
+    phas = np.array(phas)
+    fgrs = np.array(fgrs)
+
+    imageio.mimwrite(f'{output_path}/{video_name}_fgr.mp4', fgrs, fps=fps, quality=7)
+    imageio.mimwrite(f'{output_path}/{video_name}_pha.mp4', phas, fps=fps, quality=7)
+
+
+@app.command()
+def main(
+    input_path: Annotated[str, typer.Option("-i", "--input-path",
+        help="Path of the input video or frame folder.")],
+    mask_path: Annotated[str, typer.Option("-m", "--mask-path",
+        help="Path of the first-frame segmentation mask.")],
+    output_path: Annotated[str, typer.Option("-o", "--output-path",
+        help="Output folder.")] = "results/",
+    ckpt_path: Annotated[str, typer.Option("-c", "--ckpt-path",
+        help="Path of the MatAnyone2 model.")] = "pretrained_models/matanyone2.pth",
+    warmup: Annotated[int, typer.Option("-w", "--warmup",
+        help="Number of warmup iterations for the first frame alpha prediction.")] = 10,
+    erode_kernel: Annotated[int, typer.Option("-e", "--erode-kernel",
+        help="Erosion kernel size on the input mask.")] = 10,
+    dilate_kernel: Annotated[int, typer.Option("-d", "--dilate-kernel",
+        help="Dilation kernel size on the input mask.")] = 10,
+    suffix: Annotated[str, typer.Option(
+        help="Suffix to specify different target when saving.")] = "",
+    save_image: Annotated[bool, typer.Option("--save-image",
+        help="Save output frames.")] = False,
+    max_size: Annotated[int, typer.Option(
+        help="Downsamples if min(w, h) exceeds this value. -1 means no limit.")] = -1,
+):
+    """Run MatAnyone2 video matting inference."""
+    run_inference(
+        input_path=input_path,
+        mask_path=mask_path,
+        output_path=output_path,
+        ckpt_path=ckpt_path,
+        n_warmup=warmup,
+        r_erode=erode_kernel,
+        r_dilate=dilate_kernel,
+        suffix=suffix,
+        save_image=save_image,
+        max_size=max_size,
+    )
+
+
+if __name__ == '__main__':
+    app()
@@ -0,0 +1,109 @@
+import math
+import os
+import requests
+from torch.hub import download_url_to_file, get_dir
+from tqdm import tqdm
+from urllib.parse import urlparse
+
+def sizeof_fmt(size, suffix='B'):
+    """Get human readable file size.
+
+    Args:
+        size (int): File size.
+        suffix (str): Suffix. Default: 'B'.
+
+    Return:
+        str: Formated file siz.
+    """
+    for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
+        if abs(size) < 1024.0:
+            return f'{size:3.1f} {unit}{suffix}'
+        size /= 1024.0
+    return f'{size:3.1f} Y{suffix}'
+
+
+def download_file_from_google_drive(file_id, save_path):
+    """Download files from google drive.
+    Ref:
+    https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive  # noqa E501
+    Args:
+        file_id (str): File id.
+        save_path (str): Save path.
+    """
+
+    session = requests.Session()
+    URL = 'https://docs.google.com/uc?export=download'
+    params = {'id': file_id}
+
+    response = session.get(URL, params=params, stream=True)
+    token = get_confirm_token(response)
+    if token:
+        params['confirm'] = token
+        response = session.get(URL, params=params, stream=True)
+
+    # get file size
+    response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
+    print(response_file_size)
+    if 'Content-Range' in response_file_size.headers:
+        file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
+    else:
+        file_size = None
+
+    save_response_content(response, save_path, file_size)
+
+
+def get_confirm_token(response):
+    for key, value in response.cookies.items():
+        if key.startswith('download_warning'):
+            return value
+    return None
+
+
+def save_response_content(response, destination, file_size=None, chunk_size=32768):
+    if file_size is not None:
+        pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
+
+        readable_file_size = sizeof_fmt(file_size)
+    else:
+        pbar = None
+
+    with open(destination, 'wb') as f:
+        downloaded_size = 0
+        for chunk in response.iter_content(chunk_size):
+            downloaded_size += chunk_size
+            if pbar is not None:
+                pbar.update(1)
+                pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
+            if chunk:  # filter out keep-alive new chunks
+                f.write(chunk)
+        if pbar is not None:
+            pbar.close()
+
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Load file form http url, will download models if necessary.
+    Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
+    Args:
+        url (str): URL to be downloaded.
+        model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
+            Default: None.
+        progress (bool): Whether to show the download progress. Default: True.
+        file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
+    Returns:
+        str: The path to the downloaded file.
+    """
+    if model_dir is None:  # use the pytorch hub_dir
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
@@ -4,7 +4,6 @@ import random
 import numpy as np

 import torch
-import torchvision

 IMAGE_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG')
 VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.MP4', '.MOV', '.AVI')
@@ -12,8 +11,18 @@ VIDEO_EXTENSIONS = ('.mp4', '.mov', '.avi', '.MP4', '.MOV', '.AVI')
 def read_frame_from_videos(frame_root):
    if frame_root.endswith(VIDEO_EXTENSIONS):  # Video file path
        video_name = os.path.basename(frame_root)[:-4]
-        frames, _, info = torchvision.io.read_video(filename=frame_root, pts_unit='sec', output_format='TCHW') # RGB
-        fps = info['video_fps']
+        cap = cv2.VideoCapture(frame_root)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if fps <= 0:
+            fps = 24
+        frames = []
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frames.append(frame[..., [2, 1, 0]])  # BGR to RGB, HWC
+        cap.release()
+        frames = torch.Tensor(np.array(frames)).permute(0, 3, 1, 2).contiguous()  # TCHW
    else:
        video_name = os.path.basename(frame_root)
        frames = []
@@ -14,7 +14,7 @@ column_limit = 100
 name = "matanyone2"
 version = "1.0.0"
 authors = [{ name = "Peiqing Yang", email = "peiqingyang99@outlook.com" }]
-description = ""
+description = "Robust and Real-Time Video Matting with Temporal Guidance"
 readme = "README.md"
 requires-python = ">=3.10"
 classifiers = [
@@ -22,43 +22,85 @@ classifiers = [
  "Operating System :: OS Independent",
 ]
 dependencies = [
+  'numpy >= 1.21',
+  'Pillow >= 9.5',
+  'opencv-python >= 4.8',
+  'scipy >= 1.7',
+  'tqdm >= 4.66.1',
+  'einops >= 0.6',
+  'hydra-core >= 1.3.2',
+  'requests',
+  'imageio >= 2.25.0',
+  'imageio[ffmpeg]',
+  'huggingface_hub >= 0.25.0',
+  'safetensors',
+  'kornia',
+  'easydict',
+  'torch',
+  'torchvision',
+  'typer >= 0.9.0',
+  "av>=16.1.0",
+]
+
+[project.optional-dependencies]
+dev = [
  'cython',
  'gitpython >= 3.1',
  'thinplate@git+https://github.com/cheind/py-thin-plate-spline',
  'hickle >= 5.0',
  'tensorboard >= 2.11',
-  'numpy >= 1.21',
-  'Pillow >= 9.5',
-  'opencv-python >= 4.8',
-  'scipy >= 1.7',
  'pycocotools >= 2.0.7',
-  'tqdm >= 4.66.1',
-  'gradio >= 3.34',
  'gdown >= 4.7.1',
-  'einops >= 0.6',
-  'hydra-core >= 1.3.2',
-  'PySide6 >= 6.2.0',
-  'charset-normalizer >= 3.1.0',
-  'netifaces >= 0.11.0',
-  'cchardet >= 2.1.7',
-  'easydict',
-  'av >= 0.5.2',
-  'requests',
-  'pyqtdarktheme',
-  'imageio == 2.25.0',
-  'imageio[ffmpeg]',
-  'huggingface_hub == 0.36.2',
-  'safetensors',
  'xlsxwriter',
-  'kornia',
 ]
+gui = [
+  'gradio >= 6.9.0',
+  'PySide6 >= 6.2.0',
+  'pyqtdarktheme',
+]
+all = [
+  'cython',
+  'gitpython >= 3.1',
+  'thinplate@git+https://github.com/cheind/py-thin-plate-spline',
+  'hickle >= 5.0',
+  'tensorboard >= 2.11',
+  'pycocotools >= 2.0.7',
+  'gdown >= 4.7.1',
+  'xlsxwriter',
+  'gradio >= 6.9.0',
+  'PySide6 >= 6.2.0',
+  'pyqtdarktheme',
+]
+
+[project.scripts]
+matanyone2 = "matanyone2.cli:app"

 [tool.hatch.build.targets.wheel]
 packages = ["matanyone2"]

+[tool.hatch.build.targets.wheel.force-include]
+"matanyone2/config" = "matanyone2/config"
+
+[tool.hatch.build.targets.sdist]
+include = [
+  "matanyone2/",
+  "inference_matanyone2.py",
+  "README.md",
+  "LICENSE.txt",
+]
+
+[tool.uv]
+no-build-isolation-package = ["thinplate"]
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu128" }
+torchvision = { index = "pytorch-cu128" }
+
 [project.urls]
 "Homepage" = "https://github.com/pq-yang/MatAnyone2"
 "Bug Tracker" = "https://github.com/pq-yang/MatAnyone2/issues"
-
-[tool.setuptools]
-package-dir = {"" = "."}