resize_data.py

import cv2
import os
import argparse
import json
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor

def resize_mask(args):
    """
    调整单个mask文件的尺寸
    
    Args:
        args: 包含mask_path, output_path, target_size的元组
    """
    mask_path, output_path, target_size = args
    
    # 读取mask图像（灰度模式）
    img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    
    if img is None:
        return f"无法读取: {mask_path}"

    # 执行缩放，使用最邻近插值，适合Mask，不会产生中间灰色像素
    resized = cv2.resize(img, (target_size[0], target_size[1]), interpolation=cv2.INTER_NEAREST)

    # 保存
    cv2.imwrite(output_path, resized)
    return None

def process_masks(input_dir, output_dir, workers, cameras_json):
    """
    处理mask文件夹，从camera.json读取每个图片的尺寸并调整mask尺寸
    
    Args:
        input_dir: 输入mask文件夹路径
        output_dir: 输出mask文件夹路径
        workers: 并行进程数
        cameras_json: camera.json文件路径，用于读取每个图片的尺寸
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    valid_exts = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff')
    mask_files = [f for f in os.listdir(input_dir) if f.lower().endswith(valid_exts)]
    
    # 加载camera.json获取每个图片的尺寸
    img_size_map = {}
    with open(cameras_json, 'r') as f:
        cameras = json.load(f)
    
    for cam in cameras:
        img_name = os.path.basename(cam['img_name'])
        # 移除扩展名，以便匹配mask文件名
        img_base = os.path.splitext(img_name)[0]
        img_size_map[img_base] = (cam['width'], cam['height'])
    
    tasks = []
    for f in mask_files:
        # 获取目标尺寸
        f_base = os.path.splitext(f)[0]
        if f_base in img_size_map:
            size = img_size_map[f_base]
        else:
            print(f"警告: 未在camera.json中找到 {f_base} 的尺寸信息，将跳过此文件")
            continue
        
        tasks.append((
            os.path.join(input_dir, f),
            os.path.join(output_dir, f),
            size
        ))

    print(f"正在处理掩模: {input_dir} -> {output_dir}")
    with ProcessPoolExecutor(max_workers=workers) as executor:
        list(tqdm(executor.map(resize_mask, tasks), total=len(tasks)))

def main():
    parser = argparse.ArgumentParser(description="调整mask尺寸，使其与camera.json中每个图片的尺寸一致")
    parser.add_argument("--mask_in", type=str, required=True, help="原始掩模文件夹")
    parser.add_argument("--mask_out", type=str, required=True, help="缩放后掩模保存路径")
    parser.add_argument("--cameras", type=str, required=True, help="camera.json文件路径，用于读取每个图片的尺寸")
    parser.add_argument("--workers", type=int, default=8, help="进程数")

    args = parser.parse_args()

    # 只处理mask，不再处理图像
    process_masks(args.mask_in, args.mask_out, args.workers, args.cameras)

    print("\n[Done] 所有数据已同步缩放完成！")

if __name__ == "__main__":
    main()