[GAN]CelebAMask-HQ数据集处理、划分

CelebAMask数据集是2020年MaskGAN这篇文章一起提出的，这篇文章主要讲了利用用户修改后的语义分割信息来引导图像编辑。
港中文和商汤合作的，还是蛮不错的一篇文章。
文章链接：MaskGAN: Towards Diverse and Interactive Facial Image Manipulation
本篇博客只讲文章对于CelebAMask数据集补充的部分（因为博主目前更关心这部分哈哈哈啊哈！）

CelebAMask简介

• MaskGAN作者们构造一个名为CelebAMask-HQ的大规模高分辨率人脸数据集，其中包含细粒度的掩码标签（mask labels）。CelebAMask-HQ由超过30000张512×512分辨率的人脸图像组成，其中每张图像都带有19个面部组件类别的语义掩码（例如眼睛区域、鼻子区域、嘴巴区域）
• CelebAMask-HQ数据集和现存数据集的对比：

CelebAMask的特性

• CelebAMask-HQ大规模人脸语义标签数据集，根据CelebA-HQ[17]进行标记（CelebA-HQ包含来自CelebA[27]的30000张高分辨率人脸图像）。它有几个吸引人的特性：

感兴趣可以去看看CelebA-HQ和CelebA的提出
CelebA-HQ[17]Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.Progressive growing of gans for improved quality, stability,and variation. arXiv preprint arXiv:1710.10196, 2017.
CelebA[27] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.Progressive growing of gans for improved quality, stability,and variation. arXiv preprint arXiv:1710.10196, 2017.

1. 全面的注释。CelebAMask-HQ的尺寸为512×512，有19个类别，包括所有面部组件和配件的精准手工注释，如“皮肤”、“鼻子”、“眼睛”、“眉毛”、“耳朵”、“嘴巴”、“嘴唇”、“头发”、“帽子”、“眼镜”、“耳环”、“项链”、“脖子”和“布料”。
2. 标签大小选择。CelebA-HQ[17]中的图像大小为1024×1024。然而，CelebAMask-HQ选择了512×512的mask大小，因为在1024×1024的图像上标记人脸的成本非常高。此外，我们可以通过最近邻插值（nearest-neighbor interpolation）轻松地将标签从512×512扩展到1024×1024，而不会产生明显的伪影。
3. 质量控制。手工注释后，我们对每一个分割mask进行了质量检查。此外，我们要求注释人员通过几轮迭代来优化所有遮罩。（标数据真惨啊，真的能给人眼睛标瞎）
4. 特殊情况处理。对于遮挡的处理：如果面部组件被部分遮挡，我们要求注释者通过人类推断来标记组件被遮挡的部分。另一方面，我们跳过了那些完全被遮挡的组件的注释。

CelebAMask使用

• 数据预处理https://github.com/switchablenorms/CelebAMask-HQ/tree/master/face_parsing

g_mask.py生成完整mask

• 数据集提供的CelebAMask-HQ-mask-anno是每张图像各个类别（眼、鼻、配饰等）分开的mask，如下图：
  
• 首先先将CelebA-HQ-img（图像）和CelebAMask-HQ-mask-anno（未整合的label）移动到Data_preprocessing目录下。还有CelebA-HQ-to-CelebA-mapping.txt（CelebA-HQ和CelebA数据级的映射，用于训练集验证集测试集的的划分），这个txt文件记得把第一行去掉！。博主用直接软链接的方式，链过来了。
  
  
• 通过该脚本把mask值合在一起，需要注意修改文件名（估计是作者更新数据集的时候把数据集名字改了，处理mask的脚本也是很早写的，所以名字没对应上，跑还不报错，但是结果是没有数值的mask，看了一个小时才发现是名字没对上，也是醉了）
  修改+整理+注释的脚本如下：
  g_mask.py

import os
import cv2
import numpy as np
from utils import make_folder

label_list = ['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']

# 输入数据集的名称一定要检查一下，这个名称错了不会报错，但是输出的mask里面全部都是空值
folder_base = 'CelebAMask-HQ-mask-anno'
folder_save = 'CelebAMaskHQ-mask'
img_num = 30000 # 数据集一共有30000张图片

make_folder(folder_save)

for k in range(img_num):
    folder_num = k // 2000 # 该图片的分割组件存放的目录，一共有15个目录，每个目录存了2000张分割结果（包含一张图片的面部各个组件分开的分割结果）
    im_base = np.zeros((512, 512))
    for idx, label in enumerate(label_list):
        filename = os.path.join(folder_base, str(folder_num), str(k).rjust(5, '0') + '_' + label + '.png')
        if (os.path.exists(filename)):
            print (label, idx+1)
            im = cv2.imread(filename)
            im = im[:, :, 0]  # 取出图像第一个通道的值（分割图像只有一个通道，但是是部分的组件）
            im_base[im != 0] = (idx + 1) # 将该部分的值赋予一个idx+1的数值，实现分割，后期填充上颜色就变成我们看到的最终分割结果了

    filename_save = os.path.join(folder_save, str(k) + '.png')
    print (filename_save)
    cv2.imwrite(filename_save, im_base) # 保存图片

• 生成的图片就是黑的，上面存的都是1~20的数值。
  

g_color.py生成带颜色的mask

• 注意一下folder_base和folder_save就可以了，还有循环中有些小错，可能原作者忘更新了。
  g_color.py

import os
from PIL import Image
import numpy as np
from utils import make_folder
color_list = [[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]]

folder_base = 'CelebAMaskHQ-mask'
folder_save = 'CelebAMask-HQ-mask-color'
img_num = 30000

make_folder(folder_save)

for k in range(img_num):
    filename = os.path.join(folder_base, str(k) + '.png')
    im_base = np.zeros((512, 512, 3))
    if (os.path.exists(filename)):
        print(filename)
        im = Image.open(filename)
        im = np.array(im)
        for idx, color in enumerate(color_list):
            im_base[im == idx] = color # 将标签idx（单通道数据）和三通道的颜色对应上，给数据一个色彩填充
    filename_save = os.path.join(folder_save, str(k) + '.png')
    result = Image.fromarray((im_base).astype(np.uint8))
    print (filename_save)
    result.save(filename_save)

• 对应的RGB颜色可以上网查一下，比如[76, 153, 0]对应的就是下面的绿色
  
• 最终生成的带颜色的mask如下所示：
  

g_partition.py划分数据集

• 脚本如下，原作者可能是想保存一下train_list，但是忘写了哈哈哈哈啊。

import os
import shutil
import pandas as pd
from shutil import copyfile
from utils import make_folder

#### source data path
# s_label = 'CelebAMask-HQ-label'
s_label = 'CelebAMask-HQ-mask'
s_img = 'CelebA-HQ-img'

#### destination training data path
d_train_label = 'train_label'
d_train_img = 'train_img'
#### destination testing data path
d_test_label = 'test_label'
d_test_img = 'test_img'
#### val data path
d_val_label = 'val_label'
d_val_img = 'val_img'

#### make folder
make_folder(d_train_label)
make_folder(d_train_img)
make_folder(d_test_label)
make_folder(d_test_img)
make_folder(d_val_label)
make_folder(d_val_img)

#### calculate data counts in destination folder
train_count = 0
test_count = 0
val_count = 0

image_list = pd.read_csv('CelebA-HQ-to-CelebA-mapping.txt', delim_whitespace=True, header=None)
f_train = open('train_list.txt', 'w')
f_val = open('val_list.txt', 'w')
f_test = open('test_list.txt', 'w')

for idx, x in enumerate(image_list.loc[:, 1]):
    print (idx, x)
    if x >= 162771 and x < 182638:
        copyfile(os.path.join(s_label, str(idx)+'.png'), os.path.join(d_val_label, str(val_count)+'.png'))
        copyfile(os.path.join(s_img, str(idx)+'.jpg'), os.path.join(d_val_img, str(val_count)+'.jpg'))        
        val_count += 1

    elif x >= 182638:
        copyfile(os.path.join(s_label, str(idx)+'.png'), os.path.join(d_test_label, str(test_count)+'.png'))
        copyfile(os.path.join(s_img, str(idx)+'.jpg'), os.path.join(d_test_img, str(test_count)+'.jpg'))
        test_count += 1 
    else:
        copyfile(os.path.join(s_label, str(idx)+'.png'), os.path.join(d_train_label, str(train_count)+'.png'))
        copyfile(os.path.join(s_img, str(idx)+'.jpg'), os.path.join(d_train_img, str(train_count)+'.jpg'))
        train_count += 1  

print (train_count + test_count + val_count)
#### close the file
f_train.close()
f_val.close()
f_test.close()

• 主要是通过和CelebA中的映射完成的（左列，CelebA-HQ数据标号，右列CelebA中数据标号）
  
• CelebA数据集提供的数据集划分（0训练集，1验证集，2测试集）
  
• 分出来的训练图片大概26000张吧，验证集测试集各2000张左右~