# test_im = cubes[:test]
# test_sil = sils[:test]
# test_param = params[:test]
# number_testn_im = np.shape(test_im)[0]
test_im = cubes
test_sil = sils
test_param = params
number_testn_im = np.shape(test_im)[0]
# ------------------------------------------------------------------
normalize = Normalize(mean=[0.5], std=[0.5])
gray_to_rgb = Lambda(lambda x: x.repeat(3, 1, 1))
transforms = Compose([ToTensor(), normalize])
test_dataset = CubeDataset(test_im, test_sil, test_param, transforms)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=2)
# for image, sil, param in test_dataloader:
#
# nim = image.size()[0]
# for i in range(0,nim):
# print(image.size(), sil.size(), param.size()) #torch.Size([batch, 3, 512, 512]) torch.Size([batch, 6])
# im = i
# print(param[im]) # parameter in form tensor([2.5508, 0.0000, 0.0000, 0.0000, 0.0000, 5.0000])
#
#
val_im = cubes[:split] # remaining ratio for validation
val_sil = sils[:split]
val_param = params[:split]
test_im = cubes[split:split + testlen]
test_sil = sils[split:split + testlen]
test_param = params[split:split + testlen]
number_testn_im = np.shape(test_im)[0]
# ------------------------------------------------------------------
normalize = Normalize(mean=[0.5], std=[0.5])
gray_to_rgb = Lambda(lambda x: x.repeat(3, 1, 1))
transforms = Compose([ToTensor(), normalize])
train_dataset = CubeDataset(train_im, train_sil, train_param, transforms)
val_dataset = CubeDataset(val_im, val_sil, val_param, transforms)
test_dataset = CubeDataset(test_im, test_sil, test_param, transforms)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=2)
val_dataloader = DataLoader(val_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=2)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=2)
BackgroundVal = np.load(Background_Valfile)
silsVal = np.load(BWShaft_Valfile)
paramsVal = np.load(parameters_Valfile)
# print(np.min(params[:,4]))
# ------------------------------------------------------------------
val_im = BackgroundVal[start:start + vallen] #100:200
val_sil = silsVal[start:start + vallen]
val_param = paramsVal[start:start + vallen]
# ------------------------------------------------------------------
normalize = Normalize(mean=[0.5], std=[0.5])
transforms = Compose([ToTensor(), normalize])
val_dataset = CubeDataset(val_im, val_sil, val_param, transforms)
val_dataloader = DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=2)
# ------------------------------------------------------------------
# Setup the model
current_dir = os.path.dirname(os.path.realpath(__file__))
data_dir = os.path.join(current_dir, '3D_objects')
noise = 0.0
parser = argparse.ArgumentParser()
parser.add_argument('-io',
