# Create Output Paths
if not os.path.exists(output_dir):
os.makedirs(output_dir)
################################################################################
################################## Data Setup ##################################
################################################################################
# Load Data
train = data.load_images_from_folder(train_data_folder, use_noise_images)
val = data.load_images_from_folder(val_data_folder, use_noise_images)
test = data.load_images_from_folder(test_data_folder, use_noise_images)
# Preprocess & Create Data Loaders
transform = data.image2tensor_resize(image_size)
train_data = data.CustomDataset(train, transform)
val_data = data.CustomDataset(val, transform)
test_data = data.CustomDataset(test, transform)
train_dataloader = torch.utils.data.DataLoader(
train_data,
batch_size=batch_size,
shuffle=True,
num_workers=num_dataloader_workers,
pin_memory=gpu,
)
val_dataloader = torch.utils.data.DataLoader(
val_data,
batch_size=batch_size,
gpu = torch.cuda.is_available()
device = torch.device("cuda" if gpu else "cpu")
print(gpu, device)
# Create Output Paths
if not os.path.exists(reconstructed_dir):
os.makedirs(reconstructed_dir)
if not os.path.exists(generated_dir):
os.makedirs(generated_dir)
################################################################################
################################## Data Setup ##################################
################################################################################
# Preprocess & Create Data Loaders
transform = data.image2tensor_resize(vq_vae_image_size)
# Load Data
train_data = data.CustomDatasetNoMemory(train_data_folder, transform,
use_noise_images)
val_data = data.CustomDatasetNoMemory(val_data_folder, transform,
use_noise_images)
test_data = data.CustomDatasetNoMemory(test_data_folder, transform,
use_noise_images)
train_dataloader = torch.utils.data.DataLoader(
train_data,
batch_size=batch_size,
shuffle=True,
num_workers=num_dataloader_workers,
pin_memory=gpu,
