Spaces:
Runtime error
Runtime error
File size: 7,727 Bytes
c7f097c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
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)
|