You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
104 lines
3.4 KiB
104 lines
3.4 KiB
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve. |
|
# |
|
#Licensed under the Apache License, Version 2.0 (the "License"); |
|
#you may not use this file except in compliance with the License. |
|
#You may obtain a copy of the License at |
|
# |
|
# http://www.apache.org/licenses/LICENSE-2.0 |
|
# |
|
#Unless required by applicable law or agreed to in writing, software |
|
#distributed under the License is distributed on an "AS IS" BASIS, |
|
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
#See the License for the specific language governing permissions and |
|
#limitations under the License. |
|
|
|
import cv2 |
|
import os.path |
|
import numpy as np |
|
import paddle |
|
from .base_dataset import BaseDataset |
|
from .image_folder import ImageFolder |
|
|
|
from .builder import DATASETS |
|
from .preprocess.builder import build_transforms |
|
|
|
|
|
@DATASETS.register() |
|
class AnimeGANV2Dataset(paddle.io.Dataset): |
|
""" |
|
""" |
|
|
|
def __init__(self, |
|
dataroot, |
|
style, |
|
transform_real=None, |
|
transform_anime=None, |
|
transform_gray=None): |
|
"""Initialize this dataset class. |
|
|
|
Args: |
|
dataroot (dict): Directory of dataset. |
|
|
|
""" |
|
self.root = dataroot |
|
self.style = style |
|
|
|
self.transform_real = build_transforms(transform_real) |
|
self.transform_anime = build_transforms(transform_anime) |
|
self.transform_gray = build_transforms(transform_gray) |
|
|
|
self.real_root = os.path.join(self.root, 'train_photo') |
|
self.anime_root = os.path.join(self.root, f'{self.style}', 'style') |
|
self.smooth_root = os.path.join(self.root, f'{self.style}', 'smooth') |
|
|
|
self.real = ImageFolder( |
|
self.real_root, transform=self.transform_real, loader=self.loader) |
|
self.anime = ImageFolder( |
|
self.anime_root, transform=self.transform_anime, loader=self.loader) |
|
self.anime_gray = ImageFolder( |
|
self.anime_root, transform=self.transform_gray, loader=self.loader) |
|
self.smooth_gray = ImageFolder( |
|
self.smooth_root, transform=self.transform_gray, loader=self.loader) |
|
self.sizes = [ |
|
len(fold) for fold in [self.real, self.anime, self.smooth_gray] |
|
] |
|
self.size = max(self.sizes) |
|
self.reshuffle() |
|
|
|
@staticmethod |
|
def loader(path): |
|
return cv2.cvtColor( |
|
cv2.imread( |
|
path, flags=cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) |
|
|
|
def reshuffle(self): |
|
indexs = [] |
|
for cur_size in self.sizes: |
|
x = np.arange(0, cur_size) |
|
np.random.shuffle(x) |
|
if cur_size != self.size: |
|
pad_num = self.size - cur_size |
|
pad = np.random.choice(cur_size, pad_num, replace=True) |
|
x = np.concatenate((x, pad)) |
|
np.random.shuffle(x) |
|
indexs.append(x.tolist()) |
|
self.indexs = list(zip(*indexs)) |
|
|
|
def __getitem__(self, index): |
|
try: |
|
index = self.indexs.pop() |
|
except IndexError as e: |
|
self.reshuffle() |
|
index = self.indexs.pop() |
|
|
|
real_idx, anime_idx, smooth_idx = index |
|
|
|
return { |
|
'real': self.real[real_idx], |
|
'anime': self.anime[anime_idx], |
|
'anime_gray': self.anime_gray[anime_idx], |
|
'smooth_gray': self.smooth_gray[smooth_idx] |
|
} |
|
|
|
def __len__(self): |
|
return self.size
|
|
|