release evaluation code & doc
This commit is contained in:
pq-yang
@@ -0,0 +1,242 @@
|
||||
# This file is modified based on `evaluate_hr.py` from [RobustVideoMatting](https://github.com/PeterL1n/RobustVideoMatting)
|
||||
# Changes:
|
||||
# - Adapted for CRGNN dataset
|
||||
# - Supported metrics: pha_mad, pha_mse, pha_grad, pha_dtssd
|
||||
|
||||
"""
|
||||
HR (High-Resolution) evaluation. We found using numpy is very slow for high resolution, so we moved it to PyTorch using CUDA.
|
||||
|
||||
Note, the script only does evaluation. You will need to first inference yourself and save the results to disk
|
||||
Expected directory format for both prediction and ground-truth is:
|
||||
|
||||
crgnn
|
||||
├── video1
|
||||
├── pha
|
||||
├── 0000.png
|
||||
├── video2
|
||||
├── pha
|
||||
├── 0000.png
|
||||
|
||||
Prediction must have the exact file structure and file name as the ground-truth,
|
||||
except for the file extension.
|
||||
|
||||
Example usage:
|
||||
|
||||
python evaluation/eval_crgnn.py \
|
||||
--pred-dir ./data/results/crgnn \
|
||||
--true-dir ./data/crgnn/alpha
|
||||
|
||||
An excel sheet with evaluation results will be written to "PATH_TO_PREDICTIONS/crgnn.xlsx"
|
||||
"""
|
||||
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import cv2
|
||||
import kornia
|
||||
import numpy as np
|
||||
import xlsxwriter
|
||||
import torch
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class Evaluator:
|
||||
def __init__(self):
|
||||
self.parse_args()
|
||||
self.init_metrics()
|
||||
self.evaluate()
|
||||
self.write_excel()
|
||||
|
||||
def parse_args(self):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--pred-dir', type=str, required=True)
|
||||
parser.add_argument('--true-dir', type=str, required=True)
|
||||
parser.add_argument('--num-workers', type=int, default=48)
|
||||
parser.add_argument('--metrics', type=str, nargs='+', default=[
|
||||
'pha_mad', 'pha_mse', 'pha_grad', 'pha_dtssd'])
|
||||
self.args = parser.parse_args()
|
||||
|
||||
def init_metrics(self):
|
||||
self.mad = MetricMAD()
|
||||
self.mse = MetricMSE()
|
||||
self.grad = MetricGRAD()
|
||||
self.dtssd = MetricDTSSD()
|
||||
|
||||
def evaluate(self):
|
||||
tasks = []
|
||||
position = 0
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.args.num_workers) as executor:
|
||||
for video_folder in sorted(os.listdir(self.args.pred_dir)):
|
||||
pred_video_path = os.path.join(self.args.pred_dir, video_folder)
|
||||
true_video_path = os.path.join(self.args.true_dir, video_folder)
|
||||
if os.path.isdir(pred_video_path) and os.path.isdir(true_video_path):
|
||||
future = executor.submit(self.evaluate_worker, video_folder, position)
|
||||
tasks.append((video_folder, future))
|
||||
position += 1
|
||||
|
||||
self.results = [(video_folder, future.result()) for video_folder, future in tasks]
|
||||
|
||||
def write_excel(self):
|
||||
workbook = xlsxwriter.Workbook(os.path.join(self.args.pred_dir, f'{os.path.basename(self.args.pred_dir)}.xlsx'))
|
||||
summarysheet = workbook.add_worksheet('summary')
|
||||
metricsheets = [workbook.add_worksheet(metric) for metric in self.results[0][-1].keys()]
|
||||
|
||||
for i, metric in enumerate(self.results[0][-1].keys()):
|
||||
summarysheet.write(i, 0, metric)
|
||||
summarysheet.write(i, 1, f'={metric}!B2')
|
||||
|
||||
for row, (video_folder, metrics) in enumerate(self.results):
|
||||
for metricsheet, metric in zip(metricsheets, metrics.values()):
|
||||
# Write the header
|
||||
if row == 0:
|
||||
metricsheet.write(1, 0, 'Average')
|
||||
metricsheet.write(1, 1, f'=AVERAGE(C2:ZZ2)')
|
||||
for col in range(len(metric)):
|
||||
metricsheet.write(0, col + 2, col)
|
||||
colname = xlsxwriter.utility.xl_col_to_name(col + 2)
|
||||
metricsheet.write(1, col + 2, f'=AVERAGE({colname}3:{colname}9999)')
|
||||
|
||||
metricsheet.write(row + 2, 0, video_folder)
|
||||
metricsheet.write_row(row + 2, 1, metric)
|
||||
|
||||
workbook.close()
|
||||
|
||||
def evaluate_worker(self, video_folder, position):
|
||||
pred_pha_dir = os.path.join(self.args.pred_dir, video_folder, 'pha')
|
||||
true_pha_dir = os.path.join(self.args.true_dir, video_folder)
|
||||
|
||||
# Get sorted filenames from both directories
|
||||
pred_framenames = sorted([f for f in os.listdir(pred_pha_dir) if os.path.isfile(os.path.join(pred_pha_dir, f))])
|
||||
true_framenames = sorted([f for f in os.listdir(true_pha_dir) if os.path.isfile(os.path.join(true_pha_dir, f))])
|
||||
|
||||
# true-dir contains frames at every 10 frames (frame 1, 11, 21, ...)
|
||||
# pred-dir contains all frames (frame 0, 1, 2, 3, ...)
|
||||
# Match: true frame i (which is video frame i*10) corresponds to pred frame i*10
|
||||
num_true_frames = len(true_framenames)
|
||||
num_pred_frames = len(pred_framenames)
|
||||
|
||||
if num_true_frames == 0:
|
||||
print(f'Warning: {video_folder} has no true frames')
|
||||
return {metric_name : [] for metric_name in self.args.metrics}
|
||||
|
||||
print(f'Info: {video_folder} has {num_pred_frames} pred frames and {num_true_frames} true frames (every 10 frames)')
|
||||
print(f' First few pred files: {pred_framenames[:5]}')
|
||||
print(f' First few true files: {true_framenames[:5]}')
|
||||
|
||||
metrics = {metric_name : [] for metric_name in self.args.metrics}
|
||||
|
||||
pred_pha_tm1 = None
|
||||
true_pha_tm1 = None
|
||||
|
||||
for true_idx in tqdm(range(num_true_frames), desc=video_folder, position=position, dynamic_ncols=True):
|
||||
# true frame at index true_idx corresponds to video frame true_idx * 10
|
||||
# pred frame at index pred_idx = true_idx * 10
|
||||
pred_idx = true_idx * 10
|
||||
|
||||
if pred_idx >= num_pred_frames:
|
||||
print(f'Warning: {video_folder} true frame {true_idx} (video frame {pred_idx}) exceeds pred frames ({num_pred_frames})')
|
||||
break
|
||||
|
||||
pred_pha_path = os.path.join(pred_pha_dir, pred_framenames[pred_idx])
|
||||
true_pha_path = os.path.join(true_pha_dir, true_framenames[true_idx])
|
||||
|
||||
# Print the matched pair
|
||||
print(f' [{video_folder}] Pair {true_idx}: pred[{pred_idx}]="{pred_framenames[pred_idx]}" <-> true[{true_idx}]="{true_framenames[true_idx]}"')
|
||||
|
||||
pred_pha = cv2.imread(pred_pha_path, cv2.IMREAD_GRAYSCALE)
|
||||
true_pha = cv2.imread(true_pha_path, cv2.IMREAD_GRAYSCALE)
|
||||
|
||||
if pred_pha is None or true_pha is None:
|
||||
print(f'Warning: Failed to read image at true_idx {true_idx} (pred_idx {pred_idx}) in {video_folder}')
|
||||
continue
|
||||
|
||||
true_pha = torch.from_numpy(true_pha).cuda(non_blocking=True).float().div_(255)
|
||||
pred_pha = torch.from_numpy(pred_pha).cuda(non_blocking=True).float().div_(255)
|
||||
|
||||
if 'pha_mad' in self.args.metrics:
|
||||
metrics['pha_mad'].append(self.mad(pred_pha, true_pha))
|
||||
if 'pha_mse' in self.args.metrics:
|
||||
metrics['pha_mse'].append(self.mse(pred_pha, true_pha))
|
||||
if 'pha_grad' in self.args.metrics:
|
||||
metrics['pha_grad'].append(self.grad(pred_pha, true_pha))
|
||||
if 'pha_conn' in self.args.metrics:
|
||||
metrics['pha_conn'].append(self.conn(pred_pha, true_pha))
|
||||
if 'pha_dtssd' in self.args.metrics:
|
||||
if true_idx == 0:
|
||||
metrics['pha_dtssd'].append(0)
|
||||
else:
|
||||
metrics['pha_dtssd'].append(self.dtssd(pred_pha, pred_pha_tm1, true_pha, true_pha_tm1))
|
||||
|
||||
pred_pha_tm1 = pred_pha
|
||||
true_pha_tm1 = true_pha
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
class MetricMAD:
|
||||
def __call__(self, pred, true):
|
||||
return (pred - true).abs_().mean() * 1e3
|
||||
|
||||
|
||||
class MetricMSE:
|
||||
def __call__(self, pred, true):
|
||||
return ((pred - true) ** 2).mean() * 1e3
|
||||
|
||||
|
||||
class MetricGRAD:
|
||||
def __init__(self, sigma=1.4):
|
||||
self.filter_x, self.filter_y = self.gauss_filter(sigma)
|
||||
self.filter_x = torch.from_numpy(self.filter_x).unsqueeze(0).cuda()
|
||||
self.filter_y = torch.from_numpy(self.filter_y).unsqueeze(0).cuda()
|
||||
|
||||
def __call__(self, pred, true):
|
||||
true_grad = self.gauss_gradient(true)
|
||||
pred_grad = self.gauss_gradient(pred)
|
||||
return ((true_grad - pred_grad) ** 2).sum() / 1000
|
||||
|
||||
def gauss_gradient(self, img):
|
||||
img_filtered_x = kornia.filters.filter2d(img[None, None, :, :], self.filter_x, border_type='replicate')[0, 0]
|
||||
img_filtered_y = kornia.filters.filter2d(img[None, None, :, :], self.filter_y, border_type='replicate')[0, 0]
|
||||
return (img_filtered_x**2 + img_filtered_y**2).sqrt()
|
||||
|
||||
@staticmethod
|
||||
def gauss_filter(sigma, epsilon=1e-2):
|
||||
half_size = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon)))
|
||||
size = int(2 * half_size + 1)
|
||||
|
||||
# create filter in x axis
|
||||
filter_x = np.zeros((size, size))
|
||||
for i in range(size):
|
||||
for j in range(size):
|
||||
filter_x[i, j] = MetricGRAD.gaussian(i - half_size, sigma) * MetricGRAD.dgaussian(
|
||||
j - half_size, sigma)
|
||||
|
||||
# normalize filter
|
||||
norm = np.sqrt((filter_x**2).sum())
|
||||
filter_x = filter_x / norm
|
||||
filter_y = np.transpose(filter_x)
|
||||
|
||||
return filter_x, filter_y
|
||||
|
||||
@staticmethod
|
||||
def gaussian(x, sigma):
|
||||
return np.exp(-x**2 / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))
|
||||
|
||||
@staticmethod
|
||||
def dgaussian(x, sigma):
|
||||
return -x * MetricGRAD.gaussian(x, sigma) / sigma**2
|
||||
|
||||
|
||||
class MetricDTSSD:
|
||||
def __call__(self, pred_t, pred_tm1, true_t, true_tm1):
|
||||
dtSSD = ((pred_t - pred_tm1) - (true_t - true_tm1)) ** 2
|
||||
dtSSD = dtSSD.sum() / true_t.numel()
|
||||
dtSSD = dtSSD.sqrt()
|
||||
return dtSSD * 1e2
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Evaluator()
|
||||
@@ -0,0 +1,213 @@
|
||||
# This file is modified based on `evaluate_hr.py` from [RobustVideoMatting](https://github.com/PeterL1n/RobustVideoMatting)
|
||||
# Changes:
|
||||
# - Adapted for YouTubeMatte/VideoMatte dataset
|
||||
# - Supported metrics: pha_mad, pha_mse, pha_grad, pha_dtssd
|
||||
|
||||
"""
|
||||
HR (High-Resolution) evaluation. We found using numpy is very slow for high resolution, so we moved it to PyTorch using CUDA.
|
||||
|
||||
Note, the script only does evaluation. You will need to first inference yourself and save the results to disk
|
||||
Expected directory format for both prediction and ground-truth is (same for VideoMatte):
|
||||
|
||||
youtubematte_1920x1080
|
||||
├── youtubematte_motion
|
||||
├── 0000
|
||||
├── pha
|
||||
├── 0000.png
|
||||
├── youtubematte_static
|
||||
├── 0000
|
||||
├── pha
|
||||
├── 0000.png
|
||||
|
||||
Prediction must have the exact file structure and file name as the ground-truth,
|
||||
except for the file extension.
|
||||
|
||||
Example usage:
|
||||
|
||||
[YouTubeMatte]
|
||||
python evaluation/eval_hr.py \
|
||||
--pred-dir ./data/results/youtubematte_1920x1080 \
|
||||
--true-dir ./data/YouTubeMatte/youtubematte_1920x1080
|
||||
|
||||
[VideoMatte]
|
||||
python evaluation/eval_hr.py \
|
||||
--pred-dir ./data/results/videomatte_1920x1080 \
|
||||
--true-dir ./data/VideoMatte/videomatte_1920x1080
|
||||
|
||||
An excel sheet with evaluation results will be written to "PATH_TO_PREDICTIONS/xxxmatte_1920x1080.xlsx"
|
||||
"""
|
||||
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import cv2
|
||||
import kornia
|
||||
import numpy as np
|
||||
import xlsxwriter
|
||||
import torch
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class Evaluator:
|
||||
def __init__(self):
|
||||
self.parse_args()
|
||||
self.init_metrics()
|
||||
self.evaluate()
|
||||
self.write_excel()
|
||||
|
||||
def parse_args(self):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--pred-dir', type=str, required=True)
|
||||
parser.add_argument('--true-dir', type=str, required=True)
|
||||
parser.add_argument('--num-workers', type=int, default=48)
|
||||
parser.add_argument('--metrics', type=str, nargs='+', default=[
|
||||
'pha_mad', 'pha_mse', 'pha_grad', 'pha_dtssd'])
|
||||
self.args = parser.parse_args()
|
||||
|
||||
def init_metrics(self):
|
||||
self.mad = MetricMAD()
|
||||
self.mse = MetricMSE()
|
||||
self.grad = MetricGRAD()
|
||||
self.dtssd = MetricDTSSD()
|
||||
|
||||
def evaluate(self):
|
||||
tasks = []
|
||||
position = 0
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.args.num_workers) as executor:
|
||||
for dataset in sorted(os.listdir(self.args.pred_dir)):
|
||||
if os.path.isdir(os.path.join(self.args.pred_dir, dataset)):
|
||||
for clip in sorted(os.listdir(os.path.join(self.args.pred_dir, dataset))):
|
||||
future = executor.submit(self.evaluate_worker, dataset, clip, position)
|
||||
tasks.append((dataset, clip, future))
|
||||
position += 1
|
||||
|
||||
self.results = [(dataset, clip, future.result()) for dataset, clip, future in tasks]
|
||||
|
||||
def write_excel(self):
|
||||
workbook = xlsxwriter.Workbook(os.path.join(self.args.pred_dir, f'{os.path.basename(self.args.pred_dir)}.xlsx'))
|
||||
summarysheet = workbook.add_worksheet('summary')
|
||||
metricsheets = [workbook.add_worksheet(metric) for metric in self.results[0][-1].keys()]
|
||||
|
||||
for i, metric in enumerate(self.results[0][-1].keys()):
|
||||
summarysheet.write(i, 0, metric)
|
||||
summarysheet.write(i, 1, f'={metric}!B2')
|
||||
|
||||
for row, (dataset, clip, metrics) in enumerate(self.results):
|
||||
for metricsheet, metric in zip(metricsheets, metrics.values()):
|
||||
# Write the header
|
||||
if row == 0:
|
||||
metricsheet.write(1, 0, 'Average')
|
||||
metricsheet.write(1, 1, f'=AVERAGE(C2:ZZ2)')
|
||||
for col in range(len(metric)):
|
||||
metricsheet.write(0, col + 2, col)
|
||||
colname = xlsxwriter.utility.xl_col_to_name(col + 2)
|
||||
metricsheet.write(1, col + 2, f'=AVERAGE({colname}3:{colname}9999)')
|
||||
|
||||
metricsheet.write(row + 2, 0, dataset)
|
||||
metricsheet.write(row + 2, 1, clip)
|
||||
metricsheet.write_row(row + 2, 2, metric)
|
||||
|
||||
workbook.close()
|
||||
|
||||
def evaluate_worker(self, dataset, clip, position):
|
||||
framenames = sorted(os.listdir(os.path.join(self.args.pred_dir, dataset, clip, 'pha')))
|
||||
|
||||
metrics = {metric_name : [] for metric_name in self.args.metrics}
|
||||
|
||||
pred_pha_tm1 = None
|
||||
true_pha_tm1 = None
|
||||
|
||||
for i, framename in enumerate(tqdm(framenames, desc=f'{dataset} {clip}', position=position, dynamic_ncols=True)):
|
||||
true_pha = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'pha', framename[:-4]+".jpg"), cv2.IMREAD_GRAYSCALE)
|
||||
pred_pha = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE)
|
||||
|
||||
true_pha = torch.from_numpy(true_pha).cuda(non_blocking=True).float().div_(255)
|
||||
pred_pha = torch.from_numpy(pred_pha).cuda(non_blocking=True).float().div_(255)
|
||||
|
||||
if 'pha_mad' in self.args.metrics:
|
||||
metrics['pha_mad'].append(self.mad(pred_pha, true_pha))
|
||||
if 'pha_mse' in self.args.metrics:
|
||||
metrics['pha_mse'].append(self.mse(pred_pha, true_pha))
|
||||
if 'pha_grad' in self.args.metrics:
|
||||
metrics['pha_grad'].append(self.grad(pred_pha, true_pha))
|
||||
if 'pha_conn' in self.args.metrics:
|
||||
metrics['pha_conn'].append(self.conn(pred_pha, true_pha))
|
||||
if 'pha_dtssd' in self.args.metrics:
|
||||
if i == 0:
|
||||
metrics['pha_dtssd'].append(0)
|
||||
else:
|
||||
metrics['pha_dtssd'].append(self.dtssd(pred_pha, pred_pha_tm1, true_pha, true_pha_tm1))
|
||||
|
||||
pred_pha_tm1 = pred_pha
|
||||
true_pha_tm1 = true_pha
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
class MetricMAD:
|
||||
def __call__(self, pred, true):
|
||||
return (pred - true).abs_().mean() * 1e3
|
||||
|
||||
|
||||
class MetricMSE:
|
||||
def __call__(self, pred, true):
|
||||
return ((pred - true) ** 2).mean() * 1e3
|
||||
|
||||
|
||||
class MetricGRAD:
|
||||
def __init__(self, sigma=1.4):
|
||||
self.filter_x, self.filter_y = self.gauss_filter(sigma)
|
||||
self.filter_x = torch.from_numpy(self.filter_x).unsqueeze(0).cuda()
|
||||
self.filter_y = torch.from_numpy(self.filter_y).unsqueeze(0).cuda()
|
||||
|
||||
def __call__(self, pred, true):
|
||||
true_grad = self.gauss_gradient(true)
|
||||
pred_grad = self.gauss_gradient(pred)
|
||||
return ((true_grad - pred_grad) ** 2).sum() / 1000
|
||||
|
||||
def gauss_gradient(self, img):
|
||||
img_filtered_x = kornia.filters.filter2d(img[None, None, :, :], self.filter_x, border_type='replicate')[0, 0]
|
||||
img_filtered_y = kornia.filters.filter2d(img[None, None, :, :], self.filter_y, border_type='replicate')[0, 0]
|
||||
return (img_filtered_x**2 + img_filtered_y**2).sqrt()
|
||||
|
||||
@staticmethod
|
||||
def gauss_filter(sigma, epsilon=1e-2):
|
||||
half_size = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon)))
|
||||
size = int(2 * half_size + 1)
|
||||
|
||||
# create filter in x axis
|
||||
filter_x = np.zeros((size, size))
|
||||
for i in range(size):
|
||||
for j in range(size):
|
||||
filter_x[i, j] = MetricGRAD.gaussian(i - half_size, sigma) * MetricGRAD.dgaussian(
|
||||
j - half_size, sigma)
|
||||
|
||||
# normalize filter
|
||||
norm = np.sqrt((filter_x**2).sum())
|
||||
filter_x = filter_x / norm
|
||||
filter_y = np.transpose(filter_x)
|
||||
|
||||
return filter_x, filter_y
|
||||
|
||||
@staticmethod
|
||||
def gaussian(x, sigma):
|
||||
return np.exp(-x**2 / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))
|
||||
|
||||
@staticmethod
|
||||
def dgaussian(x, sigma):
|
||||
return -x * MetricGRAD.gaussian(x, sigma) / sigma**2
|
||||
|
||||
|
||||
class MetricDTSSD:
|
||||
def __call__(self, pred_t, pred_tm1, true_t, true_tm1):
|
||||
dtSSD = ((pred_t - pred_tm1) - (true_t - true_tm1)) ** 2
|
||||
dtSSD = dtSSD.sum() / true_t.numel()
|
||||
dtSSD = dtSSD.sqrt()
|
||||
return dtSSD * 1e2
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Evaluator()
|
||||
@@ -0,0 +1,252 @@
|
||||
# This file is modified based on `evaluate_lr.py` from [RobustVideoMatting](https://github.com/PeterL1n/RobustVideoMatting)
|
||||
# Changes:
|
||||
# - Adapted for YouTubeMatte/VideoMatte dataset
|
||||
# - Supported metrics: pha_mad, pha_mse, pha_grad, pha_dtssd, pha_conn
|
||||
|
||||
"""
|
||||
LR (Low-Resolution) evaluation.
|
||||
|
||||
Note, the script only does evaluation. You will need to first inference yourself and save the results to disk
|
||||
Expected directory format for both prediction and ground-truth is (same for VideoMatte):
|
||||
|
||||
youtubematte_512x288
|
||||
├── youtubematte_motion
|
||||
├── 0000
|
||||
├── pha
|
||||
├── 0000.png
|
||||
├── youtubematte_static
|
||||
├── 0000
|
||||
├── pha
|
||||
├── 0000.png
|
||||
|
||||
Prediction must have the exact file structure and file name as the ground-truth,
|
||||
meaning that if the ground-truth is png/jpg, prediction should be png/jpg.
|
||||
|
||||
Example usage:
|
||||
|
||||
[YouTubeMatte]
|
||||
python evaluation/eval_lr.py \
|
||||
--pred-dir ./data/results/youtubematte_512x288 \
|
||||
--true-dir ./data/YouTubeMatte/youtubematte_512x288
|
||||
|
||||
[VideoMatte]
|
||||
python evaluation/eval_lr.py \
|
||||
--pred-dir ./data/results/videomatte_512x288 \
|
||||
--true-dir ./data/VideoMatte/videomatte_512x288
|
||||
|
||||
An excel sheet with evaluation results will be written to "PATH_TO_PREDICTIONS/xxxmatte_512x288.xlsx"
|
||||
"""
|
||||
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
import xlsxwriter
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
class Evaluator:
|
||||
def __init__(self):
|
||||
self.parse_args()
|
||||
self.init_metrics()
|
||||
self.evaluate()
|
||||
self.write_excel()
|
||||
|
||||
def parse_args(self):
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--pred-dir', type=str, required=True)
|
||||
parser.add_argument('--true-dir', type=str, required=True)
|
||||
parser.add_argument('--num-workers', type=int, default=48)
|
||||
parser.add_argument('--metrics', type=str, nargs='+', default=[
|
||||
'pha_mad', 'pha_mse', 'pha_grad', 'pha_dtssd', 'pha_conn'])
|
||||
self.args = parser.parse_args()
|
||||
|
||||
def init_metrics(self):
|
||||
self.mad = MetricMAD()
|
||||
self.mse = MetricMSE()
|
||||
self.grad = MetricGRAD()
|
||||
self.conn = MetricCONN()
|
||||
self.dtssd = MetricDTSSD()
|
||||
|
||||
def evaluate(self):
|
||||
tasks = []
|
||||
position = 0
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.args.num_workers) as executor:
|
||||
for dataset in sorted(os.listdir(self.args.pred_dir)):
|
||||
if os.path.isdir(os.path.join(self.args.pred_dir, dataset)):
|
||||
for clip in sorted(os.listdir(os.path.join(self.args.pred_dir, dataset))):
|
||||
future = executor.submit(self.evaluate_worker, dataset, clip, position)
|
||||
tasks.append((dataset, clip, future))
|
||||
position += 1
|
||||
|
||||
self.results = [(dataset, clip, future.result()) for dataset, clip, future in tasks]
|
||||
|
||||
def write_excel(self):
|
||||
workbook = xlsxwriter.Workbook(os.path.join(self.args.pred_dir, f'{os.path.basename(self.args.pred_dir)}.xlsx'))
|
||||
summarysheet = workbook.add_worksheet('summary')
|
||||
metricsheets = [workbook.add_worksheet(metric) for metric in self.results[0][-1].keys()]
|
||||
|
||||
for i, metric in enumerate(self.results[0][-1].keys()):
|
||||
summarysheet.write(i, 0, metric)
|
||||
summarysheet.write(i, 1, f'={metric}!B2')
|
||||
|
||||
for row, (dataset, clip, metrics) in enumerate(self.results):
|
||||
for metricsheet, metric in zip(metricsheets, metrics.values()):
|
||||
# Write the header
|
||||
if row == 0:
|
||||
metricsheet.write(1, 0, 'Average')
|
||||
metricsheet.write(1, 1, f'=AVERAGE(C2:ZZ2)')
|
||||
for col in range(len(metric)):
|
||||
metricsheet.write(0, col + 2, col)
|
||||
colname = xlsxwriter.utility.xl_col_to_name(col + 2)
|
||||
metricsheet.write(1, col + 2, f'=AVERAGE({colname}3:{colname}9999)')
|
||||
|
||||
metricsheet.write(row + 2, 0, dataset)
|
||||
metricsheet.write(row + 2, 1, clip)
|
||||
metricsheet.write_row(row + 2, 2, metric)
|
||||
|
||||
workbook.close()
|
||||
|
||||
def evaluate_worker(self, dataset, clip, position):
|
||||
framenames = sorted(os.listdir(os.path.join(self.args.pred_dir, dataset, clip, 'pha')))
|
||||
|
||||
metrics = {metric_name : [] for metric_name in self.args.metrics}
|
||||
|
||||
pred_pha_tm1 = None
|
||||
true_pha_tm1 = None
|
||||
|
||||
for i, framename in enumerate(tqdm(framenames, desc=f'{dataset} {clip}', position=position, dynamic_ncols=True)):
|
||||
true_pha = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'pha', framename[:-4]+".png"), cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255
|
||||
pred_pha = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255
|
||||
|
||||
if 'pha_mad' in self.args.metrics:
|
||||
metrics['pha_mad'].append(self.mad(pred_pha, true_pha))
|
||||
if 'pha_mse' in self.args.metrics:
|
||||
metrics['pha_mse'].append(self.mse(pred_pha, true_pha))
|
||||
if 'pha_grad' in self.args.metrics:
|
||||
metrics['pha_grad'].append(self.grad(pred_pha, true_pha))
|
||||
if 'pha_conn' in self.args.metrics:
|
||||
metrics['pha_conn'].append(self.conn(pred_pha, true_pha))
|
||||
if 'pha_dtssd' in self.args.metrics:
|
||||
if i == 0:
|
||||
metrics['pha_dtssd'].append(0)
|
||||
else:
|
||||
metrics['pha_dtssd'].append(self.dtssd(pred_pha, pred_pha_tm1, true_pha, true_pha_tm1))
|
||||
|
||||
pred_pha_tm1 = pred_pha
|
||||
true_pha_tm1 = true_pha
|
||||
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
class MetricMAD:
|
||||
def __call__(self, pred, true):
|
||||
return np.abs(pred - true).mean() * 1e3
|
||||
|
||||
|
||||
class MetricMSE:
|
||||
def __call__(self, pred, true):
|
||||
return ((pred - true) ** 2).mean() * 1e3
|
||||
|
||||
|
||||
class MetricGRAD:
|
||||
def __init__(self, sigma=1.4):
|
||||
self.filter_x, self.filter_y = self.gauss_filter(sigma)
|
||||
|
||||
def __call__(self, pred, true):
|
||||
pred_normed = np.zeros_like(pred)
|
||||
true_normed = np.zeros_like(true)
|
||||
cv2.normalize(pred, pred_normed, 1., 0., cv2.NORM_MINMAX)
|
||||
cv2.normalize(true, true_normed, 1., 0., cv2.NORM_MINMAX)
|
||||
|
||||
true_grad = self.gauss_gradient(true_normed).astype(np.float32)
|
||||
pred_grad = self.gauss_gradient(pred_normed).astype(np.float32)
|
||||
|
||||
grad_loss = ((true_grad - pred_grad) ** 2).sum()
|
||||
return grad_loss / 1000
|
||||
|
||||
def gauss_gradient(self, img):
|
||||
img_filtered_x = cv2.filter2D(img, -1, self.filter_x, borderType=cv2.BORDER_REPLICATE)
|
||||
img_filtered_y = cv2.filter2D(img, -1, self.filter_y, borderType=cv2.BORDER_REPLICATE)
|
||||
return np.sqrt(img_filtered_x**2 + img_filtered_y**2)
|
||||
|
||||
@staticmethod
|
||||
def gauss_filter(sigma, epsilon=1e-2):
|
||||
half_size = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon)))
|
||||
size = int(2 * half_size + 1)
|
||||
|
||||
# create filter in x axis
|
||||
filter_x = np.zeros((size, size))
|
||||
for i in range(size):
|
||||
for j in range(size):
|
||||
filter_x[i, j] = MetricGRAD.gaussian(i - half_size, sigma) * MetricGRAD.dgaussian(
|
||||
j - half_size, sigma)
|
||||
|
||||
# normalize filter
|
||||
norm = np.sqrt((filter_x**2).sum())
|
||||
filter_x = filter_x / norm
|
||||
filter_y = np.transpose(filter_x)
|
||||
|
||||
return filter_x, filter_y
|
||||
|
||||
@staticmethod
|
||||
def gaussian(x, sigma):
|
||||
return np.exp(-x**2 / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))
|
||||
|
||||
@staticmethod
|
||||
def dgaussian(x, sigma):
|
||||
return -x * MetricGRAD.gaussian(x, sigma) / sigma**2
|
||||
|
||||
|
||||
class MetricCONN:
|
||||
def __call__(self, pred, true):
|
||||
step=0.1
|
||||
thresh_steps = np.arange(0, 1 + step, step)
|
||||
round_down_map = -np.ones_like(true)
|
||||
for i in range(1, len(thresh_steps)):
|
||||
true_thresh = true >= thresh_steps[i]
|
||||
pred_thresh = pred >= thresh_steps[i]
|
||||
intersection = (true_thresh & pred_thresh).astype(np.uint8)
|
||||
|
||||
# connected components
|
||||
_, output, stats, _ = cv2.connectedComponentsWithStats(
|
||||
intersection, connectivity=4)
|
||||
# start from 1 in dim 0 to exclude background
|
||||
size = stats[1:, -1]
|
||||
|
||||
# largest connected component of the intersection
|
||||
omega = np.zeros_like(true)
|
||||
if len(size) != 0:
|
||||
max_id = np.argmax(size)
|
||||
# plus one to include background
|
||||
omega[output == max_id + 1] = 1
|
||||
|
||||
mask = (round_down_map == -1) & (omega == 0)
|
||||
round_down_map[mask] = thresh_steps[i - 1]
|
||||
round_down_map[round_down_map == -1] = 1
|
||||
|
||||
true_diff = true - round_down_map
|
||||
pred_diff = pred - round_down_map
|
||||
# only calculate difference larger than or equal to 0.15
|
||||
true_phi = 1 - true_diff * (true_diff >= 0.15)
|
||||
pred_phi = 1 - pred_diff * (pred_diff >= 0.15)
|
||||
|
||||
connectivity_error = np.sum(np.abs(true_phi - pred_phi))
|
||||
return connectivity_error / 1000
|
||||
|
||||
|
||||
class MetricDTSSD:
|
||||
def __call__(self, pred_t, pred_tm1, true_t, true_tm1):
|
||||
dtSSD = ((pred_t - pred_tm1) - (true_t - true_tm1)) ** 2
|
||||
dtSSD = np.sum(dtSSD) / true_t.size
|
||||
dtSSD = np.sqrt(dtSSD)
|
||||
return dtSSD * 1e2
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Evaluator()
|
||||
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
|
||||
input_folder="data/crgnn/image_allframe"
|
||||
mask_folder="data/crgnn/mask"
|
||||
|
||||
ckpt_name="matanyone2"
|
||||
|
||||
for video_folder in "${input_folder}"/*; do
|
||||
if [ -d "${video_folder}" ]; then
|
||||
video_id=$(basename "${video_folder}")
|
||||
|
||||
mask_file="${mask_folder}/${video_id}.png"
|
||||
if [ -f "${mask_file}" ]; then
|
||||
|
||||
input_frames_folder="${video_folder}"
|
||||
if [ -d "${input_frames_folder}" ]; then
|
||||
echo "Processing video: ${video_id}"
|
||||
|
||||
python inference_matanyone2.py \
|
||||
--input_path "${input_frames_folder}" \
|
||||
--mask_path "${mask_file}" \
|
||||
--output_path "data/results/crgnn" \
|
||||
--ckpt_path "pretrained_models/${ckpt_name}.pth" \
|
||||
--warmup 10 \
|
||||
--erode_kernel 10 \
|
||||
--dilate_kernel 10 \
|
||||
--save_image
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
done
|
||||
@@ -0,0 +1,38 @@
|
||||
#!/bin/bash
|
||||
|
||||
input_folder="./data/VideoMatte/videomatte_1920x1080"
|
||||
mask_folder="./data/VideoMatte_first_frame_seg_mask/videomatte_1920x1080"
|
||||
|
||||
ckpt_name="matanyone2"
|
||||
|
||||
for subfolder in "videomatte_motion" "videomatte_static"; do
|
||||
subfolder_path="${input_folder}/${subfolder}"
|
||||
|
||||
echo "Processing subfolder: ${subfolder}"
|
||||
|
||||
for video_folder in "${subfolder_path}"/*; do
|
||||
if [ -d "${video_folder}" ]; then
|
||||
video_id=$(basename "${video_folder}")
|
||||
|
||||
mask_file="${mask_folder}/${subfolder}/${video_id}.png"
|
||||
if [ -f "${mask_file}" ]; then
|
||||
|
||||
input_frames_folder="${video_folder}/com"
|
||||
if [ -d "${input_frames_folder}" ]; then
|
||||
echo "Processing video: ${video_id} from ${subfolder}"
|
||||
|
||||
python evaluation/inference_matanyone_eval.py \
|
||||
--input_path "${input_frames_folder}" \
|
||||
--mask_path "${mask_file}" \
|
||||
--output_path "./data/results/videomatte_1920x1080/${subfolder}" \
|
||||
--ckpt_path "pretrained_models/${ckpt_name}.pth" \
|
||||
--warmup 10 \
|
||||
--erode_kernel 15 \
|
||||
--dilate_kernel 15 \
|
||||
--save_image
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
#!/bin/bash
|
||||
|
||||
input_folder="./data/YouTubeMatte/youtubematte_1920x1080"
|
||||
mask_folder="./data/YouTubeMatte_first_frame_seg_mask/youtubematte_1920x1080"
|
||||
|
||||
ckpt_name="matanyone2"
|
||||
|
||||
for subfolder in "youtubematte_motion" "youtubematte_static"; do
|
||||
subfolder_path="${input_folder}/${subfolder}"
|
||||
|
||||
echo "Processing subfolder: ${subfolder}"
|
||||
|
||||
for video_folder in "${subfolder_path}"/*; do
|
||||
if [ -d "${video_folder}" ]; then
|
||||
video_id=$(basename "${video_folder}")
|
||||
|
||||
mask_file="${mask_folder}/${video_id}.png"
|
||||
if [ -f "${mask_file}" ]; then
|
||||
|
||||
input_frames_folder="${video_folder}/har"
|
||||
if [ -d "${input_frames_folder}" ]; then
|
||||
echo "Processing video: ${video_id} from ${subfolder}"
|
||||
|
||||
python evaluation/inference_matanyone_eval.py \
|
||||
--input_path "${input_frames_folder}" \
|
||||
--mask_path "${mask_file}" \
|
||||
--output_path "./data/results/youtubematte_1920x1080/${subfolder}" \
|
||||
--ckpt_path "pretrained_models/${ckpt_name}.pth" \
|
||||
--warmup 10 \
|
||||
--erode_kernel 15 \
|
||||
--dilate_kernel 15 \
|
||||
--save_image
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
done
|
||||
done
|
||||
@@ -0,0 +1,38 @@
|
||||
#!/bin/bash
|
||||
|
||||
input_folder="./data/VideoMatte/videomatte_512x288"
|
||||
mask_folder="./data/VideoMatte_first_frame_seg_mask/videomatte_512x288"
|
||||
|
||||
ckpt_name="matanyone2"
|
||||
|
||||
for subfolder in "videomatte_motion" "videomatte_static"; do
|
||||
subfolder_path="${input_folder}/${subfolder}"
|
||||
|
||||
echo "Processing subfolder: ${subfolder}"
|
||||
|
||||
for video_folder in "${subfolder_path}"/*; do
|
||||
if [ -d "${video_folder}" ]; then
|
||||
video_id=$(basename "${video_folder}")
|
||||
|
||||
mask_file="${mask_folder}/${subfolder}/${video_id}.png"
|
||||
if [ -f "${mask_file}" ]; then
|
||||
|
||||
input_frames_folder="${video_folder}/com"
|
||||
if [ -d "${input_frames_folder}" ]; then
|
||||
echo "Processing video: ${video_id} from ${subfolder}"
|
||||
|
||||
python evaluation/inference_matanyone_eval.py \
|
||||
--input_path "${input_frames_folder}" \
|
||||
--mask_path "${mask_file}" \
|
||||
--output_path "./data/results/videomatte_512x288/${subfolder}" \
|
||||
--ckpt_path "pretrained_models/${ckpt_name}.pth" \
|
||||
--warmup 1 \
|
||||
--erode_kernel 4 \
|
||||
--dilate_kernel 4 \
|
||||
--save_image
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
done
|
||||
done
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
#!/bin/bash
|
||||
|
||||
input_folder="./data/YouTubeMatte/youtubematte_512x288"
|
||||
mask_folder="./data/YouTubeMatte_first_frame_seg_mask/youtubematte_512x288"
|
||||
|
||||
ckpt_name="matanyone2"
|
||||
|
||||
for subfolder in "youtubematte_motion" "youtubematte_static"; do
|
||||
subfolder_path="${input_folder}/${subfolder}"
|
||||
|
||||
echo "Processing subfolder: ${subfolder}"
|
||||
|
||||
for video_folder in "${subfolder_path}"/*; do
|
||||
if [ -d "${video_folder}" ]; then
|
||||
video_id=$(basename "${video_folder}")
|
||||
|
||||
mask_file="${mask_folder}/${video_id}.png"
|
||||
if [ -f "${mask_file}" ]; then
|
||||
|
||||
input_frames_folder="${video_folder}/har"
|
||||
if [ -d "${input_frames_folder}" ]; then
|
||||
echo "Processing video: ${video_id} from ${subfolder}"
|
||||
|
||||
python evaluation/inference_matanyone_eval.py \
|
||||
--input_path "${input_frames_folder}" \
|
||||
--mask_path "${mask_file}" \
|
||||
--output_path "./data/results/youtubematte_512x288/${subfolder}" \
|
||||
--ckpt_path "pretrained_models/${ckpt_name}.pth" \
|
||||
--warmup 1 \
|
||||
--erode_kernel 4 \
|
||||
--dilate_kernel 4 \
|
||||
--save_image
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
done
|
||||
done
|
||||
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
import cv2
|
||||
import tqdm
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
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
|
||||
from matanyone2.utils.inference_utils import gen_dilate, gen_erosion, read_frame_from_videos
|
||||
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
device = get_default_device()
|
||||
|
||||
@torch.inference_mode()
|
||||
@safe_autocast_decorator()
|
||||
def main(input_path, mask_path, output_path, ckpt_path, n_warmup=10, r_erode=10, r_dilate=10, suffix="", save_image=False, max_size=-1):
|
||||
|
||||
video_name = os.path.basename(os.path.dirname(input_path))
|
||||
|
||||
# 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, _ = read_frame_from_videos(input_path)
|
||||
repeated_frames = vframes[0].unsqueeze(0).repeat(n_warmup, 1, 1, 1) # repeat the first frame for warmup
|
||||
vframes = torch.cat([repeated_frames, vframes], dim=0).float()
|
||||
length += n_warmup # update length
|
||||
|
||||
# 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")
|
||||
|
||||
# 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)) # green screen to paste fgr
|
||||
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: # resize needed
|
||||
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)):
|
||||
# load the image as RGB; normalization is done within the model
|
||||
image = vframes[ti]
|
||||
|
||||
image_np = np.array(image.permute(1,2,0)) # for output visualize
|
||||
image = (image / 255.).float().to(device) # for network input
|
||||
|
||||
if ti == 0:
|
||||
output_prob = processor.step(image, mask, objects=objects) # encode given mask
|
||||
output_prob = processor.step(image, first_frame_pred=True) # first frame for prediction
|
||||
else:
|
||||
if ti <= n_warmup:
|
||||
output_prob = processor.step(image, first_frame_pred=True) # reinit as the first frame for prediction
|
||||
else:
|
||||
output_prob = processor.step(image)
|
||||
|
||||
# convert output probabilities to alpha matte
|
||||
mask = processor.output_prob_to_mask(output_prob)
|
||||
|
||||
# visualize prediction
|
||||
pha = mask.unsqueeze(2).cpu().numpy()
|
||||
com_np = image_np / 255. * pha + bgr * (1 - pha)
|
||||
|
||||
# DONOT save the warmup frame
|
||||
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}/pha/{str(ti-n_warmup).zfill(4)}.png', pha)
|
||||
cv2.imwrite(f'{output_path}/{video_name}/fgr/{str(ti-n_warmup).zfill(4)}.png', com_np[...,[2,1,0]])
|
||||
|
||||
# [optional] save videos for better visualization
|
||||
# import imageio
|
||||
|
||||
# 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)
|
||||
|
||||
if __name__ == '__main__':
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-i', '--input_path', type=str, default="inputs/video/test-sample1.mp4", help='Path of the input video or frame folder.')
|
||||
parser.add_argument('-m', '--mask_path', type=str, default="inputs/mask/test-sample1.png", help='Path of the first-frame segmentation mask.')
|
||||
parser.add_argument('-o', '--output_path', type=str, default="results/", help='Output folder. Default: results')
|
||||
parser.add_argument('-c', '--ckpt_path', type=str, default="pretrained_models/matanyone.pth", help='Path of the MatAnyone model.')
|
||||
parser.add_argument('-w', '--warmup', type=str, default="10", help='Number of warmup iterations for the first frame alpha prediction.')
|
||||
parser.add_argument('-e', '--erode_kernel', type=str, default="10", help='Erosion kernel on the input mask.')
|
||||
parser.add_argument('-d', '--dilate_kernel', type=str, default="10", help='Dilation kernel on the input mask.')
|
||||
parser.add_argument('--suffix', type=str, default="", help='Suffix to specify different target when saving, e.g., target1.')
|
||||
parser.add_argument('--save_image', action='store_true', default=False, help='Save output frames. Default: False')
|
||||
parser.add_argument('--max_size', type=str, default="-1", help='When positive, the video will be downsampled if min(w, h) exceeds. Default: -1 (means no limit)')
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
main(input_path=args.input_path, \
|
||||
mask_path=args.mask_path, \
|
||||
output_path=args.output_path, \
|
||||
ckpt_path=args.ckpt_path, \
|
||||
n_warmup=args.warmup, \
|
||||
r_erode=args.erode_kernel, \
|
||||
r_dilate=args.dilate_kernel, \
|
||||
suffix=args.suffix, \
|
||||
save_image=args.save_image, \
|
||||
max_size=args.max_size)
|
||||
Reference in New Issue
Block a user