diff --git a/README.md b/README.md index 7a68a2d..a948af2 100644 --- a/README.md +++ b/README.md @@ -49,13 +49,13 @@ ## 📮 Update -- [2026.03] Release inference codes and gradio demo. +- [2026.03] Release inference codes, evaluation codes, and gradio demo. - [2025.12] This repo is created. ## 🏄🏻‍♀️ TODO - [x] Release inference codes and gradio demo. -- [ ] Release evaluation codes. +- [x] Release evaluation codes. - [ ] Release training codes for video matting model. - [ ] Release checkpoint and training codes for quality evaluator model. - [ ] Release real-world video matting dataset **VMReal**. @@ -138,6 +138,8 @@ By launching, an interactive interface will appear as follow. ![overall_teaser](assets/teaser_demo.gif) +## 📊 Evaluation +Please refer to the [evaluation documentation](docs/EVAL.md) for details. ## 🛠️ Data Pipeline ![data_pipeline](assets/data_pipeline.jpg) diff --git a/docs/EVAL.md b/docs/EVAL.md new file mode 100644 index 0000000..224df78 --- /dev/null +++ b/docs/EVAL.md @@ -0,0 +1,153 @@ +# Evaluation Documentation + +This document shows the way to reproduce the quantitative results in our paper on two synthetic benchmarks ([YouTubeMatte](https://github.com/pq-yang/MatAnyone/tree/main?tab=readme-ov-file#youtubematte-dataset) and [VideoMatte](https://github.com/PeterL1n/RobustVideoMatting/blob/master/documentation/training.md#evaluation)) and one real benchmark ([CRGNN](https://github.com/TiantianWang/VideoMatting-CRGNN)). + +## YouTubeMatte + +**📦 We provide the inference results with MatAnyone 2 on the YouTubeMatte benchmark [here](https://drive.google.com/drive/folders/1fgPAx4pRGyxGIYW4NeBevDM8PpA0TG9I?usp=sharing).** + +### Preparation +* [YouTubeMatte.zip (6.24G)](https://drive.google.com/file/d/1IEH0RaimT_hSp38AWF6wuwNJzzNSHpJ4/view?usp=drive_link) +* [YouTubeMatte_first_frame_seg_mask.zip (310K)](https://drive.google.com/file/d/1Zpa7SB7VZmkvRDiehVC-c_0dmFWXdfzK/view?usp=drive_linkk) + +To run the inference scripts, your files should be arranged as: +``` +data + |- YouTubeMatte_first_frame_seg_mask # for inference only + |- YouTubeMatte + |- youtubematte_512x288 + |- youtubematte_1920x1080 +``` + +### Batch Inference +Empirically, for low-resolution (`youtubematte_512x288`) and high-resolution (`youtubematte_1920x1080`) data, we set **different** hyperparameter values for `--warmup`, `--erode_kernel`, and `--dilate_kernel`. + +```shell +# lr: youtubematte_512x288 +bash evaluation/infer_batch_lr_yt.sh + +# hr: youtubematte_1920x1080 +bash evaluation/infer_batch_hr_yt.sh +``` + +### Evaluation +To run the evaluation scripts, your files should be arranged as: + +``` +data + |- YouTubeMatte + |- youtubematte_512x288 + |- youtubematte_1920x1080 + + |- results + |- youtubematte_512x288 + |- youtubematte_1920x1080 +``` + +```shell +# lr: youtubematte_512x288 +python evaluation/eval_lr.py \ + --pred-dir ./data/results/youtubematte_512x288 \ + --true-dir ./data/YouTubeMatte/youtubematte_512x288 + +# hr: youtubematte_1920x1080 +python evaluation/eval_hr.py \ + --pred-dir ./data/results/youtubematte_1920x1080 \ + --true-dir ./data/YouTubeMatte/youtubematte_1920x1080 +``` + +## VideoMatte + +**📦 We provide the inference results with MatAnyone 2 on the VideoMatte benchmark [here](https://drive.google.com/drive/folders/12QM-_kyerE1tQfINoR17zYIFtrcwCK04?usp=sharing).** + +### Preparation +* [videomatte_512x512.tar (1.8G)](https://robustvideomatting.blob.core.windows.net/eval/videomatte_512x288.tar) +* [videomatte_1920x1080.tar (2.2G)](https://robustvideomatting.blob.core.windows.net/eval/videomatte_1920x1080.tar) +* [VideoMatte_first_frame_seg_mask.zip (416K)](https://drive.google.com/file/d/1kN5gX4NAEa4HG-k2ir8kPcEp_18DbDHt/view?usp=drive_link) + +To run the inference scripts, your files should be arranged as: +``` +data + |- VideoMatte_first_frame_seg_mask # for inference only + |- VideoMatte + |- videomatte_512x288 + |- videomatte_1920x1080 +``` + +### Batch Inference +Empirically, for low-resolution (`videomatte_512x288`) and high-resolution (`videomatte_1920x1080`) data, we set **different** hyperparameter values for `--warmup`, `--erode_kernel`, and `--dilate_kernel`. + +```shell +# lr: videomatte_512x288 +bash evaluation/infer_batch_lr_vm.sh + +# hr: videomatte_1920x1080 +bash evaluation/infer_batch_hr_vm.sh +``` + +### Evaluation +To run the evaluation scripts, your files should be arranged as: + +``` +data + |- VideoMatte + |- videomatte_512x288 + |- videomatte_1920x1080 + + |- results + |- videomatte_512x288 + |- videomatte_512x288 +``` + +```shell +# lr: videomatte_512x288 +python evaluation/eval_lr.py \ + --pred-dir ./data/results/videomatte_512x288 \ + --true-dir ./data/VideoMatte/videomatte_512x288 + +# hr: videomatte_1920x1080 +python evaluation/eval_hr.py \ + --pred-dir ./data/results/videomatte_1920x1080 \ + --true-dir ./data/VideoMatte/videomatte_1920x1080 +``` + +## CRGNN + +**📦 We provide the inference results with MatAnyone 2 on the CRGNN benchmark [here](https://drive.google.com/file/d/1JJyE4uPymEcijNa1Ok8ME_BY6oxJ3EXJ/view?usp=sharing).** + +### Preparation +* [real_human_data](https://www.dropbox.com/sh/23uvsue5we7e7b5/AAB4GSSWIaKiSouvN3wuWiwWa?dl=0) +* [CRGNN_first_frame_seg_mask.zip (151K)](https://drive.google.com/file/d/1cDSf1kO_tdWy-q3CuX4IfuZ16ZYEER-d/view?usp=sharing) + +To run the inference scripts, your files should be arranged as: +``` +data + |- crgnn + |- alpha + |- image_allframe + |- mask # first frame seg mask +``` + +### Batch Inference + +```shell +bash evaluation/infer_batch_crgnn.sh +``` + +### Evaluation +To run the evaluation scripts, your files should be arranged as: + +``` +data + |- crgnn + |- alpha + + |- results + |- crgnn +``` + +```shell +python evaluation/eval_crgnn.py \ + --pred-dir ./data/results/crgnn \ + --true-dir ./data/crgnn/alpha +``` diff --git a/evaluation/eval_crgnn.py b/evaluation/eval_crgnn.py new file mode 100644 index 0000000..3797b0f --- /dev/null +++ b/evaluation/eval_crgnn.py @@ -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() \ No newline at end of file diff --git a/evaluation/eval_hr.py b/evaluation/eval_hr.py new file mode 100644 index 0000000..37a3e46 --- /dev/null +++ b/evaluation/eval_hr.py @@ -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() \ No newline at end of file diff --git a/evaluation/eval_lr.py b/evaluation/eval_lr.py new file mode 100644 index 0000000..6531c0a --- /dev/null +++ b/evaluation/eval_lr.py @@ -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() \ No newline at end of file diff --git a/evaluation/infer_batch_crgnn.sh b/evaluation/infer_batch_crgnn.sh new file mode 100644 index 0000000..d5ce3b7 --- /dev/null +++ b/evaluation/infer_batch_crgnn.sh @@ -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 diff --git a/evaluation/infer_batch_hr_vm.sh b/evaluation/infer_batch_hr_vm.sh new file mode 100644 index 0000000..9d58e0a --- /dev/null +++ b/evaluation/infer_batch_hr_vm.sh @@ -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 + diff --git a/evaluation/infer_batch_hr_yt.sh b/evaluation/infer_batch_hr_yt.sh new file mode 100644 index 0000000..abf43b0 --- /dev/null +++ b/evaluation/infer_batch_hr_yt.sh @@ -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 \ No newline at end of file diff --git a/evaluation/infer_batch_lr_vm.sh b/evaluation/infer_batch_lr_vm.sh new file mode 100644 index 0000000..806dba4 --- /dev/null +++ b/evaluation/infer_batch_lr_vm.sh @@ -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 + diff --git a/evaluation/infer_batch_lr_yt.sh b/evaluation/infer_batch_lr_yt.sh new file mode 100644 index 0000000..65c52b9 --- /dev/null +++ b/evaluation/infer_batch_lr_yt.sh @@ -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 \ No newline at end of file diff --git a/evaluation/inference_matanyone_eval.py b/evaluation/inference_matanyone_eval.py new file mode 100644 index 0000000..b115bac --- /dev/null +++ b/evaluation/inference_matanyone_eval.py @@ -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)