Edit2Perceive/utils/eval_matting.py

from utils.metric import compute_matting_metrics
from PIL import Image
import numpy as np
import os
import argparse
import torch
import torch.nn.functional as F
from tqdm import tqdm # 使用tqdm来显示进度条，对并行任务尤其友好
import multiprocessing as mp # 核心库

# Helper functions (resize_tensor, load_image_rgb_or_grayscale) 保持不变
def resize_tensor(input_tensor, target_height, target_width):
    """
    使用双线性插值调整深度图像大小
    """
    input_tensor = torch.from_numpy(input_tensor).unsqueeze(0).unsqueeze(0).float()
    resized_tensor = F.interpolate(
        input_tensor, 
        size=(target_height, target_width), 
        mode='bilinear', 
        align_corners=True
    )
    return resized_tensor.squeeze().numpy()

def load_image_rgb_or_grayscale(image_path):
    """
    加载图像，支持RGB和灰度图像，统一转换为numpy数组
    """
    try:
        if image_path.endswith('.npy'):
            img_array = np.load(image_path)
        else:
            img = Image.open(image_path)
            img_array = np.array(img)
            if len(img_array.shape) == 3 and img_array.shape[2] == 4:
                img_array = img_array[:, :, :3]
        return np.mean(img_array, axis=2) if len(img_array.shape) == 3 else img_array
    except Exception as e:
        # 增加错误日志，方便调试
        print(f"Warning: Failed to load image from {image_path}. Error: {e}")
        return None

# ------------------- 新增：独立的工作函数 -------------------
# 这个函数包含了处理单个样本的所有逻辑，它将在一个独立的进程中被调用。
def process_sample(task_args):
    """
    处理单个样本：加载数据、预处理、计算指标。
    Args:
        task_args (tuple): 包含单个任务所需所有路径的元组 (pred_path, alpha_path, trimap_path)
    Returns:
        tuple: 包含 (mse, mad, sad, grad, conn) 的元组，如果出错则返回五个 np.nan
    """
    pred_path, alpha_path, trimap_path = task_args

    try:
        # 1. 加载数据
        pred_depth = load_image_rgb_or_grayscale(pred_path)
        alpha = load_image_rgb_or_grayscale(alpha_path)
        trimap = load_image_rgb_or_grayscale(trimap_path)

        # 如果任何一个文件加载失败，则跳过
        if pred_depth is None or alpha is None or trimap is None:
            return (np.nan, np.nan, np.nan, np.nan, np.nan)

        # 2. 预处理 (与原eval函数中的逻辑相同)
        gt_depth = alpha
        gt_depth[np.isinf(gt_depth)] = 0
        gt_depth[np.isnan(gt_depth)] = 0
        pred_depth[np.isinf(pred_depth)] = 0
        pred_depth[np.isnan(pred_depth)] = 0

        if pred_depth.shape != gt_depth.shape:
            pred_depth = resize_tensor(pred_depth, gt_depth.shape[0], gt_depth.shape[1])
        
        # pred_depth = pred_depth > 0.5
        gt_depth = gt_depth.astype(np.float32) / 255.0

        # 3. 计算指标
        mse, mad, sad, grad, conn = compute_matting_metrics(pred_depth, gt_depth, trimap, whole=True)
        
        # 检查结果是否有效
        if not all(v >= 0 for v in [mse, mad, sad, grad, conn]):
             return (np.nan, np.nan, np.nan, np.nan, np.nan)

        return (mse, mad, sad, grad, conn)

    except Exception as e:
        print(f'Error processing {os.path.basename(pred_path)}: {e}')
        # 返回NaN，方便后续统计时忽略错误样本
        return (np.nan, np.nan, np.nan, np.nan, np.nan)

# ------------------- 重构后的 test 函数 (接口不变) -------------------
def test(args):
    # 1. 准备任务列表 (与原代码相同)
    # 这部分逻辑不变，只是现在我们不加载图像，而是收集文件路径来创建任务
    dataset_paths = {
        "comp": "./data_split/comp_matting/filenames_test.txt",
        "aim": "./data_split/AIM_matting/filenames_val.txt",
        "p3m": "./data_split/P3M_matting/filenames_val_P.txt",
        "p3m-np": "./data_split/P3M_matting/filenames_val_NP.txt",
        "am": "./data_split/AM_matting/filenames_val.txt",
    }
    
    filenames_path = dataset_paths.get(args.dataset)
    if not filenames_path:
        raise ValueError(f"Unknown dataset: {args.dataset}")

    with open(filenames_path, "r") as f:
        test_filenames = [line.strip() for line in f.readlines()]
    
    # 在非full模式下截断文件名列表
    if "full" not in args.pred_path and "result" not in args.pred_path and args.dataset in ["comp", "p3m", "p3m-np", "am"]:
        test_filenames = test_filenames[:10]

    if getattr(args, 'max_samples', None) is not None:
        test_filenames = test_filenames[:args.max_samples]

    # 2. 创建任务列表，每个任务是一个包含所需文件路径的元组
    tasks = []
    for line in test_filenames:
        items = line.split()
        # 兼容不同格式的 filenames.txt
        merged_path, trimap_path, alpha_path = items[0], items[1], items[-1]
        
        pred_full_path = os.path.join(args.pred_path, merged_path.replace(".jpg",".npy").replace(".png",".npy"))
        alpha_full_path = os.path.join(args.gt_path, alpha_path)
        trimap_full_path = os.path.join(args.gt_path, trimap_path)
        
        tasks.append((pred_full_path, alpha_full_path, trimap_full_path))

    print(f'### Starting evaluation for {len(tasks)} files on dataset "{args.dataset}"...')
    mp_context = mp.get_context('spawn')
    # 3. 使用多进程池执行任务
    # os.cpu_count() 可以获取机器的CPU核心数，自动利用所有资源
    # `with`语句可以确保进程池在使用后被正确关闭
    results_list = []
    with mp_context.Pool(processes=32) as pool:
        # pool.imap_unordered 比 map 更高效，因为它不保证返回结果的顺序
        # 使用tqdm来包装，可以实时看到处理进度
        results_iterator = pool.imap_unordered(process_sample, tasks)
        
        for result in tqdm(results_iterator, total=len(tasks), desc="Processing samples"):
            results_list.append(result)

    # 4. 汇总结果
    # 将结果列表转换为numpy数组，方便进行列操作
    results_array = np.array(results_list)
    
    # 使用 np.nanmean 来计算平均值，它会自动忽略所有NaN值（即处理失败的样本）
    # 这比手动过滤更简洁、更安全
    mean_mse = np.nanmean(results_array[:, 0])
    mean_mad = np.nanmean(results_array[:, 1])
    mean_sad = np.nanmean(results_array[:, 2])
    mean_grad = np.nanmean(results_array[:, 3])
    mean_conn = np.nanmean(results_array[:, 4])

    # 统计有效结果的数量
    valid_results_count = np.sum(~np.isnan(results_array[:, 0]))

    # 5. 打印结果 (与原代码格式一致)
    final_results_str = "{:.4f}, {:.4f}, {:.4f}, {:.4f}, {:.4f}".format(
        mean_mse, mean_mad, mean_sad, mean_grad, mean_conn
    )
    
    print('Done.')
    print(f'Valid results: {valid_results_count}/{len(tasks)}')
    print("{:>7}, {:>7}, {:>7}, {:>7}, {:>7}".format('MSE', 'MAD', 'SAD', 'Grad', 'Conn'))
    print(final_results_str)
    
    return final_results_str

# 主程序入口保持不变
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='BTS TensorFlow implementation.', fromfile_prefix_chars='@')
    parser.add_argument('--pred_path',           type=str,   help='path to the prediction results in png')
    parser.add_argument('--gt_path',             type=str,   help='root path to the groundtruth data' )
    parser.add_argument('--dataset',             type=str,   help='dataset to test on, e.g., p3m, comp, etc.', default='p3m')
    parser.add_argument('--max_samples',        type=int,    help='maximum number of samples to test', default=None)
    args = parser.parse_args()
    test(args)