class Celeba:
image_size = 256
batch_size = 8
transform = albumentations.Compose([
albumentations.HorizontalFlip(),
albumentations.Resize(CelebaWithKeyPoints.image_size, CelebaWithKeyPoints.image_size),
# albumentations.ElasticTransform(p=0.5, alpha=50, alpha_affine=1, sigma=10),
albumentations.ShiftScaleRotate(p=0.5, rotate_limit=10, scale_limit=(-0.1, 0.3)),
albumentations.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
AlbToTensor()
])
def __init__(self):
print("init calaba")
dataset = ImageDataset(
f"{Paths.default.data()}/celeba",
img_transform=Celeba.transform
)
print("dataset size: ", len(dataset))
self.loader = data.DataLoader(
dataset,
batch_size=Celeba.batch_size,
sampler=data_sampler(dataset, shuffle=True, distributed=False),
drop_last=True,
num_workers=Celeba.batch_size
)
print("batch size: ", Celeba.batch_size)
self.loader = sample_data(self.loader)
def transform():
return albumentations.Compose([
albumentations.HorizontalFlip(),
albumentations.Resize(CelebaWithKeyPoints.image_size, CelebaWithKeyPoints.image_size),
# albumentations.ElasticTransform(p=0.5, alpha=100, alpha_affine=1, sigma=10),
albumentations.ShiftScaleRotate(p=0.5, rotate_limit=10),
albumentations.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
AlbToTensor()
])
def __init__(self):
self.landmarks_path = os.path.join(Paths.default.data(), "celeba_landmarks")
self.img_names = [img[:-4] for img in os.listdir(os.path.join(self.landmarks_path))]
self.img_path = os.path.join(Paths.default.data(), "celeba/img_align_celeba/")
self.img_transform = albumentations.Compose([
albumentations.Resize(256, 256),
albumentations.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
AlbToTensor()
])
channel_multiplier=args.channel_multiplier)) #.to(device)
g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
image_size = args.size
transform = albumentations.Compose([
albumentations.HorizontalFlip(),
albumentations.Resize(image_size, image_size),
albumentations.ElasticTransform(p=0.5,
alpha=100,
alpha_affine=1,
sigma=10),
albumentations.ShiftScaleRotate(p=0.5, rotate_limit=10),
albumentations.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
AlbToTensor()
])
# dataset = ImageFolder("/raid/data/celeba", transform=transform)
dataset = ImageDataset(f"{Paths.default.data()}/celeba",
img_transform=transform)
loader = data.DataLoader(dataset,
batch_size=args.batch,
sampler=data_sampler(dataset,
shuffle=True,
distributed=False),
drop_last=True,
num_workers=20)