def predict():
#nonlocal output_dict
#model.eval()
with torch.no_grad():
input_dict_cuda, label_dict_cuda = utils_data.nestedDictToDevice((input_dict, label_dict), device=device)
output_dict_cuda = model(input_dict_cuda)
output_dict = utils_data.nestedDictToDevice(output_dict_cuda, device='cpu')
def _inference(engine, batch):
# now compute error
model.eval()
with torch.no_grad():
x, y = utils_data.nestedDictToDevice(batch, device=device) # make it work for dict input too
y_pred = model(x)
return y_pred, y
def _update(engine, batch):
model.train()
optimizer.zero_grad()
x, y = utils_data.nestedDictToDevice(batch, device=device) # make it work for dict input too
y_pred = model(x)
loss = loss_fn(y_pred, y)
loss.backward()
optimizer.step()
return loss.item(), y_pred
input_dict['bg_crop'] = input_dict['img_crop']
# white background
input_dict['bg_crop'] = torch.from_numpy(np.ones((2,3,128,128))*3).float().cuda()
input_dict['shuffled_appearance'] =torch.from_numpy(np.array([1,0]).astype('int'))
# test on example
# input_dict = ex
input_dict['external_rotation_global'] = torch.from_numpy(np.eye(3)).float().cuda()
#input_dict['external_rotation_global'] = torch.from_numpy(rotationMatrixXZY(theta=0, phi=0, psi=2.0)).float().cuda()
input_dict['external_rotation_cam'] = torch.from_numpy(np.eye(3)).float().cuda()
output_dict = None
label_dict = None
model.eval()
with torch.no_grad():
input_dict_cuda = utils_data.nestedDictToDevice((input_dict), device=device)
output_dict_cuda = model(input_dict_cuda)
output_dict = utils_data.nestedDictToDevice(output_dict_cuda, device='cpu')
skimage.io.imsave('whiteresults/' + filename + '_' + str(k) + '.png', tensor_to_img(output_dict['img_crop'][0]))
#skimage.io.imsave('redin.png', X[i,:,:,:])
#skimage.io.imsave('results/' + str(i) + '.png', Y[i,:,:,:])
#pickle
#pred = model.predict(X[:-2])
#print(Y[0])
#skimage.io.imsave('test1.png', Y[0])
#print(np.shape(Y[0]))
#testimg = tensor_to_npimg(npimg_to_tensor(Y[0]))
#skimage.io.imsave('test.png', testimg)
#break
#break