import torch
import numpy as np
from .mesh_util import *
from .sample_util import *
from .geometry import *
import cv2
from PIL import Image
from tqdm import tqdm
def reshape_multiview_tensors(image_tensor, calib_tensor):
# Careful here! Because we put single view and multiview together,
# the returned tensor.shape is 5-dim: [B, num_views, C, W, H]
# So we need to convert it back to 4-dim [B*num_views, C, W, H]
# Don't worry classifier will handle multi-view cases
image_tensor = image_tensor.view(
image_tensor.shape[0] * image_tensor.shape[1],
image_tensor.shape[2],
image_tensor.shape[3],
image_tensor.shape[4]
)
calib_tensor = calib_tensor.view(
calib_tensor.shape[0] * calib_tensor.shape[1],
calib_tensor.shape[2],
calib_tensor.shape[3]
)
return image_tensor, calib_tensor
def reshape_sample_tensor(sample_tensor, num_views):
if num_views == 1:
return sample_tensor
# Need to repeat sample_tensor along the batch dim num_views times
sample_tensor = sample_tensor.unsqueeze(dim=1)
sample_tensor = sample_tensor.repeat(1, num_views, 1, 1)
sample_tensor = sample_tensor.view(
sample_tensor.shape[0] * sample_tensor.shape[1],
sample_tensor.shape[2],
sample_tensor.shape[3]
)
return sample_tensor
def gen_mesh(opt, net, cuda, data, save_path, use_octree=True):
image_tensor = data['img'].to(device=cuda)
calib_tensor = data['calib'].to(device=cuda)
net.filter(image_tensor)
b_min = data['b_min']
b_max = data['b_max']
try:
save_img_path = save_path[:-4] + '.png'
save_img_list = []
for v in range(image_tensor.shape[0]):
save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0
save_img_list.append(save_img)
save_img = np.concatenate(save_img_list, axis=1)
Image.fromarray(np.uint8(save_img[:,:,::-1])).save(save_img_path)
verts, faces, _, _ = reconstruction(
net, cuda, calib_tensor, opt.resolution, b_min, b_max, use_octree=use_octree)
verts_tensor = torch.from_numpy(verts.T).unsqueeze(0).to(device=cuda).float()
xyz_tensor = net.projection(verts_tensor, calib_tensor[:1])
uv = xyz_tensor[:, :2, :]
color = index(image_tensor[:1], uv).detach().cpu().numpy()[0].T
color = color * 0.5 + 0.5
save_obj_mesh_with_color(save_path, verts, faces, color)
except Exception as e:
print(e)
print('Can not create marching cubes at this time.')
def gen_mesh_color(opt, netG, netC, cuda, data, save_path, use_octree=True):
image_tensor = data['img'].to(device=cuda)
calib_tensor = data['calib'].to(device=cuda)
netG.filter(image_tensor)
netC.filter(image_tensor)
netC.attach(netG.get_im_feat())
b_min = data['b_min']
b_max = data['b_max']
try:
save_img_path = save_path[:-4] + '.png'
save_img_list = []
for v in range(image_tensor.shape[0]):
save_img = (np.transpose(image_tensor[v].detach().cpu().numpy(), (1, 2, 0)) * 0.5 + 0.5)[:, :, ::-1] * 255.0
save_img_list.append(save_img)
save_img = np.concatenate(save_img_list, axis=1)
Image.fromarray(np.uint8(save_img[:,:,::-1])).save(save_img_path)
verts, faces, _, _ = reconstruction(
netG, cuda, calib_tensor, opt.resolution, b_min, b_max, use_octree=use_octree)
# Now Getting colors
verts_tensor = torch.from_numpy(verts.T).unsqueeze(0).to(device=cuda).float()
verts_tensor = reshape_sample_tensor(verts_tensor, opt.num_views)
color = np.zeros(verts.shape)
interval = 10000
for i in range(len(color) // interval):
left = i * interval
right = i * interval + interval
if i == len(color) // interval - 1:
right = -1
netC.query(verts_tensor[:, :, left:right], calib_tensor)
rgb = netC.get_preds()[0].detach().cpu().numpy() * 0.5 + 0.5
color[left:right] = rgb.T
save_obj_mesh_with_color(save_path, verts, faces, color)
except Exception as e:
print(e)
print('Can not create marching cubes at this time.')
def adjust_learning_rate(optimizer, epoch, lr, schedule, gamma):
"""Sets the learning rate to the initial LR decayed by schedule"""
if epoch in schedule:
lr *= gamma
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
def compute_acc(pred, gt, thresh=0.5):
'''
return:
IOU, precision, and recall
'''
with torch.no_grad():
vol_pred = pred > thresh
vol_gt = gt > thresh
union = vol_pred | vol_gt
inter = vol_pred & vol_gt
true_pos = inter.sum().float()
union = union.sum().float()
if union == 0:
union = 1
vol_pred = vol_pred.sum().float()
if vol_pred == 0:
vol_pred = 1
vol_gt = vol_gt.sum().float()
if vol_gt == 0:
vol_gt = 1
return true_pos / union, true_pos / vol_pred, true_pos / vol_gt
def calc_error(opt, net, cuda, dataset, num_tests):
if num_tests > len(dataset):
num_tests = len(dataset)
with torch.no_grad():
erorr_arr, IOU_arr, prec_arr, recall_arr = [], [], [], []
for idx in tqdm(range(num_tests)):
data = dataset[idx * len(dataset) // num_tests]
# retrieve the data
image_tensor = data['img'].to(device=cuda)
calib_tensor = data['calib'].to(device=cuda)
sample_tensor = data['samples'].to(device=cuda).unsqueeze(0)
if opt.num_views > 1:
sample_tensor = reshape_sample_tensor(sample_tensor, opt.num_views)
label_tensor = data['labels'].to(device=cuda).unsqueeze(0)
res, error = net.forward(image_tensor, sample_tensor, calib_tensor, labels=label_tensor)
IOU, prec, recall = compute_acc(res, label_tensor)
# print(
# '{0}/{1} | Error: {2:06f} IOU: {3:06f} prec: {4:06f} recall: {5:06f}'
# .format(idx, num_tests, error.item(), IOU.item(), prec.item(), recall.item()))
erorr_arr.append(error.item())
IOU_arr.append(IOU.item())
prec_arr.append(prec.item())
recall_arr.append(recall.item())
return np.average(erorr_arr), np.average(IOU_arr), np.average(prec_arr), np.average(recall_arr)
def calc_error_color(opt, netG, netC, cuda, dataset, num_tests):
if num_tests > len(dataset):
num_tests = len(dataset)
with torch.no_grad():
error_color_arr = []
for idx in tqdm(range(num_tests)):
data = dataset[idx * len(dataset) // num_tests]
# retrieve the data
image_tensor = data['img'].to(device=cuda)
calib_tensor = data['calib'].to(device=cuda)
color_sample_tensor = data['color_samples'].to(device=cuda).unsqueeze(0)
if opt.num_views > 1:
color_sample_tensor = reshape_sample_tensor(color_sample_tensor, opt.num_views)
rgb_tensor = data['rgbs'].to(device=cuda).unsqueeze(0)
netG.filter(image_tensor)
_, errorC = netC.forward(image_tensor, netG.get_im_feat(), color_sample_tensor, calib_tensor, labels=rgb_tensor)
# print('{0}/{1} | Error inout: {2:06f} | Error color: {3:06f}'
# .format(idx, num_tests, errorG.item(), errorC.item()))
error_color_arr.append(errorC.item())
return np.average(error_color_arr)