data/dataset.py

import os
import pickle
import random
from pathlib import Path
from typing import Dict
from typing import List

import torch
from loguru import logger
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from configs.mode import FaceSwapMode
from configs.train_config import TrainConfig


class ManyToManyTrainDataset(Dataset):
    def __init__(self, dataset_root: str, dataset_index: str, same_rate=0.5):
        """
        Many-to-many 训练数据集构建
        Parameters:
        -----------
        dataset_root: str, 数据集根目录
        dataset_index: str, 数据集index文件路径
        same_rate: float, 每个batch里面相同人脸所占的比例
        """
        super(ManyToManyTrainDataset, self).__init__()
        self.transform = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.CenterCrop((256, 256)),
                transforms.ToTensor(),
            ]
        )
        self.data_root = Path(dataset_root)
        with open(dataset_index, "rb") as f:
            self.file_index = pickle.load(f, encoding="bytes")

        self.same_rate = same_rate

        self.id_list: List[str] = list(self.file_index.keys())

        # 所有id都遍历一遍，视为一个epoch
        self.length = len(self.id_list)
        self.image_num = sum([len(v) for v in self.file_index.values()])

        self.mask_dir = "mask" if TrainConfig().mouth_mask else "mask_no_mouth"
        logger.info(f"dataset contains {self.length} ids and {self.image_num} images")
        logger.info(f"will use mask mode: {self.mask_dir}")

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        source_id_index = index
        source_file = random.choice(self.file_index[self.id_list[source_id_index]])
        if random.random() < self.same_rate:
            # 在相同id的文件列表中选择
            target_file = random.choice(self.file_index[self.id_list[source_id_index]])
            same = torch.ones(1)
        else:
            # 在不同id的文件列表中选择
            target_id_index = random.choice(list(set(range(self.length)) - set([source_id_index])))
            target_file = random.choice(self.file_index[self.id_list[target_id_index]])
            same = torch.zeros(1)

        source_file = self.data_root / Path(source_file)
        target_file = self.data_root / Path(target_file)
        target_mask_file = target_file.parent.parent.parent / self.mask_dir / target_file.parent.stem / target_file.name

        target_img = Image.open(target_file.as_posix()).convert("RGB")
        source_img = Image.open(source_file.as_posix()).convert("RGB")

        target_mask = Image.open(target_mask_file.as_posix()).convert("RGB")

        source_img = self.transform(source_img)
        target_img = self.transform(target_img)
        target_mask = self.transform(target_mask)[0, :, :].unsqueeze(0)

        return {
            "source_image": source_img,
            "target_image": target_img,
            "target_mask": target_mask,
            "same": same,
            # "source_img_name": source_file.as_posix(),
            # "target_img_name": target_file.as_posix(),
            # "target_mask_name": target_mask_file.as_posix(),
        }


class OneToManyTrainDataset(Dataset):
    def __init__(self, dataset_root: str, dataset_index: str, source_name: str, same_rate=0.5):
        """
        One-to-many 训练数据集构建
        Parameters:
        -----------
        dataset_root: str, 数据集根目录
        dataset_index: str, 数据集index文件路径
        source_name: str, source face id的名称, one-to-many里面的one
        same_rate: float, 每个batch里面相同人脸所占的比例
        """
        super(OneToManyTrainDataset, self).__init__()
        self.transform = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.CenterCrop((256, 256)),
                transforms.ToTensor(),
            ]
        )
        self.data_root = Path(dataset_root)
        with open(dataset_index, "rb") as f:
            self.file_index = pickle.load(f, encoding="bytes")
        self.same_rate = same_rate
        self.source_name = source_name

        self.id_list: List[str] = list(self.file_index.keys())

        try:
            self.source_id_index: int = self.id_list.index(self.source_name)
        except Exception:
            raise Exception(f"{self.source_name} not in dataset dir")

        # 所有id都遍历一遍，视为一个epoch
        self.length = len(self.id_list)
        self.image_num = sum([len(v) for v in self.file_index.values()])
        self.mask_dir = "mask" if TrainConfig().mouth_mask else "mask_no_mouth"
        logger.info(f"dataset contains {self.length} ids and {self.image_num} images")
        logger.info(f"will use mask mode: {self.mask_dir}")

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        target_id_index = index
        target_file = random.choice(self.file_index[self.id_list[target_id_index]])
        if random.random() < self.same_rate:
            # 在相同id的文件列表中选择
            source_file = random.choice(self.file_index[self.id_list[target_id_index]])
            same = torch.ones(1)
        else:
            # 直接选择source name中的图片
            source_file = random.choice(self.file_index[self.source_name])
            # 如果和target同个id
            if self.source_id_index == target_id_index:
                same = torch.ones(1)
            else:
                same = torch.zeros(1)

        source_file = self.data_root / Path(source_file)
        target_file = self.data_root / Path(target_file)
        target_mask_file = target_file.parent.parent.parent / self.mask_dir / target_file.parent.stem / target_file.name

        target_img = Image.open(target_file.as_posix()).convert("RGB")
        source_img = Image.open(source_file.as_posix()).convert("RGB")

        target_mask = Image.open(target_mask_file.as_posix()).convert("RGB")

        source_img = self.transform(source_img)
        target_img = self.transform(target_img)
        target_mask = self.transform(target_mask)[0, :, :].unsqueeze(0)

        return {
            "source_image": source_img,
            "target_image": target_img,
            "target_mask": target_mask,
            "same": same,
            # "source_img_name": source_file.as_posix(),
            # "target_img_name": target_file.as_posix(),
            # "target_mask_name": target_mask_file.as_posix(),
        }


class TrainDatasetDataLoader:
    """Wrapper class of Dataset class that performs multi-threaded data loading"""

    def __init__(self):
        """Initialize this class"""
        opt = TrainConfig()
        if opt.mode is FaceSwapMode.MANY_TO_MANY:
            self.dataset = ManyToManyTrainDataset(opt.dataset_root, opt.dataset_index, opt.same_rate)
        elif opt.mode is FaceSwapMode.ONE_TO_MANY:
            logger.info(f"In one-to-many mode, source face is {opt.source_name}")
            self.dataset = OneToManyTrainDataset(opt.dataset_root, opt.dataset_index, opt.source_name, opt.same_rate)
        else:
            raise NotImplementedError
        logger.info(f"dataset {type(self.dataset).__name__} created")
        if opt.use_ddp:
            self.train_sampler = torch.utils.data.distributed.DistributedSampler(self.dataset, shuffle=True)
            self.dataloader = torch.utils.data.DataLoader(
                self.dataset,
                batch_size=opt.batch_size,
                num_workers=int(opt.num_threads),
                drop_last=True,
                sampler=self.train_sampler,
                pin_memory=True,
            )
        else:
            self.dataloader = torch.utils.data.DataLoader(
                self.dataset,
                batch_size=opt.batch_size,
                shuffle=True,
                num_workers=int(opt.num_threads),
                drop_last=True,
                pin_memory=True,
            )

    def load_data(self):
        return self

    def __len__(self):
        """Return the number of data in the dataset"""
        return len(self.dataset)

    def __iter__(self):
        """Return a batch of data"""
        for data in self.dataloader:
            yield data


if __name__ == "__main__":
    dataloader = TrainDatasetDataLoader()
    for idx, data in enumerate(dataloader):
        # print(data["source_img_name"])
        # print(data["target_img_name"])
        # print(data["target_mask_name"])
        print(data["same"])