action='store_true',
help="Use normal information from the input images?")
parser.add_argument('--samples',
type=int,
default=1000,
help='Number of samples for the train and test dataset')
parser.add_argument('--batchSize',
type=int,
default=64,
help='training batch size')
opt = parser.parse_args()
device = torch.device("cpu")
# load samples and create dataset
dataset_data = datasetVideo.collect_samples_clouds_video(
opt.upscale_factor, opt)
class HighResDatasetFromSamples(torch.utils.data.Dataset):
def __init__(self):
super().__init__()
self.samples = dataset_data.samples
self.data = dataset_data
self.num_images = len(self.samples)
def get_high_res_shape(self):
return (self.data.output_channels, self.data.crop_size * 4,
self.data.crop_size * 4)
def __getitem__(self, index):
import torch
import numpy as np
import matplotlib.pyplot as plt
import datasetVideo
from utils import ScreenSpaceShading
opt = dict({'samples': 4, 'numberOfImages': 1})
dataset_data = datasetVideo.collect_samples_clouds_video(4,
opt,
deferred_shading=True)
test_full_set = datasetVideo.DatasetFromFullImages(dataset_data, 1)
input_data, input_flow = test_full_set[0]
print(input_data.shape)
device = input_data.device
shading = ScreenSpaceShading(device)
shading.fov(30)
shading.ambient_light_color(np.array([0.1, 0.1, 0.1]))
shading.diffuse_light_color(np.array([1.0, 1.0, 1.0]))
shading.specular_light_color(np.array([0.2, 0.2, 0.2]))
shading.specular_exponent(16)
shading.light_direction(np.array([0.2, 0.2, 1.0]))
shading.material_color(np.array([1.0, 0.3, 0.3]))
output_rgb = shading(input_data)
f, axarr = plt.subplots(1, 3)
axarr[0].imshow(input_data.numpy()[0, 0, :, :] * 0.5 + 0.5) #mask
axarr[1].imshow((input_data[0, 1:4, :, :] * 0.5 + 0.5).numpy().transpose(
(1, 2, 0))) #normal
print('Current run: %05d'%nextRunNumber)
runName = 'run%05d'%nextRunNumber
logdir = os.path.join(opt.logdir, runName)
modeldir = os.path.join(opt.modeldir, runName)
runName = 'run%05d'%nextRunNumber
os.makedirs(logdir)
os.makedirs(modeldir)
#########################
# DATASETS + CHANNELS
#########################
print('===> Loading datasets')
if opt.dataset.lower() == 'cloud-video':
dataset_data = datasetVideo.collect_samples_clouds_video(
opt_dict['upscale_factor'], opt_dict, deferred_shading=True)
else:
raise ValueError('Unknown dataset %s'%opt.dataset)
print('Dataset input images have %d channels'%dataset_data.input_channels)
input_channels = dataset_data.input_channels
assert input_channels == 5 # mask, normalX, normalY, normalZ, depth
output_channels = dataset_data.output_channels
assert output_channels == 6 # mask, normalX, normalY, normalZ, depth, ambient occlusion
input_channels_with_previous = input_channels + output_channels * (opt.upscale_factor ** 2)
train_set = datasetVideo.DatasetFromSamples(dataset_data, False, opt.testFraction)
test_set = datasetVideo.DatasetFromSamples(dataset_data, True, opt.testFraction)
test_full_set = datasetVideo.DatasetFromFullImages(dataset_data, min(opt.testNumFullImages, len(dataset_data.images_low)))
training_data_loader = DataLoader(dataset=train_set, batch_size=opt.batchSize, shuffle=True)
testing_data_loader = DataLoader(dataset=test_set, batch_size=opt.testBatchSize, shuffle=False)