HF integration +uv + cli
This commit is contained in:
not-lain
@@ -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 = []
|
||||
|
||||
Reference in New Issue
Block a user