保姆级教程：用Python脚本搞定CelebAMask-HQ数据集预处理与可视化（附完整代码）

从零掌握CelebAMask-HQ数据集完整预处理与可视化实战指南当第一次打开CelebAMask-HQ数据集时面对分散在多个文件夹中的19类面部组件mask文件即使是经验丰富的开发者也可能感到无从下手。这个包含3万张高分辨率人脸图像的数据集每张图都精确标注了皮肤、鼻子、眼睛等19个面部区域的语义分割mask是进行人脸编辑、虚拟化妆等计算机视觉任务的黄金标准。但原始数据组织方式让直接使用变得困难——我们需要将这些零散的标注文件整合成结构化数据并生成直观的可视化效果。1. 环境准备与数据整理在开始处理前我们需要确保Python环境已安装必要的库。建议使用conda创建专属环境conda create -n celebamask python3.8 conda activate celebamask pip install opencv-python pillow numpy pandas数据集目录结构应该如下安排CelebAMask-HQ/ ├── CelebA-HQ-img/ # 原始图像 ├── CelebAMask-HQ-mask-anno/ # 原始标注 ├── CelebA-HQ-to-CelebA-mapping.txt └── Data_preprocessing/ # 我们的处理脚本关键提示从GitHub克隆仓库时注意检查CelebAMask-HQ-mask-anno文件夹是否完整。常见问题是部分mask文件下载失败导致后续处理报错。2. Mask整合从碎片到完整标注原始数据将每个面部组件的mask单独保存我们需要将它们合并为每张图一个完整mask。创建g_mask.pyimport os import cv2 import numpy as np # 19个面部组件的定义顺序很重要影响后续可视化颜色映射 LABELS [ skin, nose, eye_g, l_eye, r_eye, l_brow, r_brow, l_ear, r_ear, mouth, u_lip, l_lip, hair, hat, ear_r, neck_l, neck, cloth ] def merge_masks(base_dir, output_dir, img_count30000): os.makedirs(output_dir, exist_okTrue) for img_id in range(img_count): # 每2000张图一个子文件夹 subfolder str(img_id // 2000) base_mask np.zeros((512, 512), dtypenp.uint8) for idx, label in enumerate(LABELS): mask_path os.path.join( base_dir, subfolder, f{img_id:05d}_{label}.png ) if os.path.exists(mask_path): mask cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) base_mask[mask 0] idx 1 # 为每个组件分配唯一编号 cv2.imwrite(f{output_dir}/{img_id}.png, base_mask)运行后会生成包含完整mask的CelebAMaskHQ-mask目录。常见问题排查报错NoneType object has no attribute shape原因mask文件缺失或路径错误解决检查CelebAMask-HQ-mask-anno下的文件名是否匹配00000_skin.png格式3. 彩色可视化让标注一目了然黑白mask不便于直观检查我们通过g_color.py添加颜色from PIL import Image import numpy as np # 为每个面部组件定义醒目颜色 COLOR_MAP [ [0,0,0], # 背景 [204,0,0], # 皮肤-红 [76,153,0], # 鼻子-绿 [204,204,0], # 眼周-黄 [51,51,255], # 左眼-蓝 [204,0,204], # 右眼-紫 [0,255,255], # 左眉-青 [255,204,204], # 右眉-粉 [102,51,0], # 左耳-棕 [255,0,0], # 右耳-亮红 [102,204,0], # 嘴-黄绿 [255,255,0], # 上唇-亮黄 [0,0,153], # 下唇-深蓝 [0,0,204], # 头发-蓝 [255,51,153], # 帽子-粉红 [0,204,204], # 耳环-青绿 [0,51,0], # 颈部-深绿 [255,153,51], # 衣领-橙 [0,204,0] # 衣服-亮绿 ] def colorize_masks(mask_dir, output_dir): os.makedirs(output_dir, exist_okTrue) for mask_file in os.listdir(mask_dir): mask_path os.path.join(mask_dir, mask_file) colored np.zeros((512,512,3), dtypenp.uint8) mask np.array(Image.open(mask_path)) for idx, color in enumerate(COLOR_MAP): colored[mask idx] color Image.fromarray(colored).save(f{output_dir}/{mask_file})生成的效果图中不同面部区域会呈现鲜明对比色。调试技巧如果某些区域颜色异常检查COLOR_MAP顺序是否与LABELS完全对应使用PIL.Image.show()快速预览结果4. 数据集划分与CelebA保持一致为确保与现有研究可比性我们按CelebA的原始划分方式拆分数据。创建g_partition.pyimport pandas as pd from shutil import copyfile def split_dataset(mask_dir, img_dir, mapping_file): # 创建输出目录 splits [train, val, test] for split in splits: os.makedirs(f{split}_label, exist_okTrue) os.makedirs(f{split}_img, exist_okTrue) # 读取CelebA的原始划分映射 mapping pd.read_csv(mapping_file, delim_whitespaceTrue, headerNone) # 按CelebA的划分标准进行分配 for idx, celebA_id in enumerate(mapping[1]): if 162771 celebA_id 182638: # 验证集 dest val elif celebA_id 182638: # 测试集 dest test else: # 训练集 dest train # 复制图像和标注 copyfile(f{mask_dir}/{idx}.png, f{dest}_label/{idx}.png) copyfile(f{img_dir}/{idx}.jpg, f{dest}_img/{idx}.jpg)最终得到的数据集结构train_img/ # 约24,000张训练图像 train_label/ # 对应的训练标注 val_img/ # 约3,000张验证图像 val_label/ test_img/ # 约3,000张测试图像 test_label/5. 高级技巧与问题解决方案在实际处理过程中有几个关键点需要特别注意文件名不匹配问题原始脚本中可能存在路径硬编码问题。现代Python推荐使用pathlib进行路径操作from pathlib import Path mask_path Path(base_dir) / subfolder / f{img_id:05d}_{label}.png if mask_path.exists(): mask cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)批量处理加速对于3万张图像可以使用多进程加速from multiprocessing import Pool def process_image(args): img_id, base_dir, output_dir args # ...处理逻辑... with Pool(8) as p: # 使用8个进程 p.map(process_image, [(i, base_dir, output_dir) for i in range(30000)])可视化检查工具创建Jupyter notebook快速检查任意图像的处理结果import matplotlib.pyplot as plt def show_sample(img_id): fig, (ax1, ax2) plt.subplots(1, 2, figsize(12,6)) img plt.imread(fCelebA-HQ-img/{img_id}.jpg) mask plt.imread(fCelebAMask-HQ-mask-color/{img_id}.png) ax1.imshow(img) ax2.imshow(mask) plt.show()这套预处理流程不仅适用于CelebAMask-HQ其设计思路也可迁移到其他分割数据集的处理中。关键在于理解数据组织逻辑、建立可靠的错误检查机制以及提供直观的可视化验证手段。