Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is | |
| # holder of all proprietary rights on this computer program. | |
| # You can only use this computer program if you have closed | |
| # a license agreement with MPG or you get the right to use the computer | |
| # program from someone who is authorized to grant you that right. | |
| # Any use of the computer program without a valid license is prohibited and | |
| # liable to prosecution. | |
| # | |
| # Copyright©2019 Max-Planck-Gesellschaft zur Förderung | |
| # der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute | |
| # for Intelligent Systems. All rights reserved. | |
| # | |
| # Contact: ps-license@tuebingen.mpg.de | |
| import yaml | |
| import os.path as osp | |
| import torch | |
| import numpy as np | |
| import torch.nn.functional as F | |
| from ..dataset.mesh_util import * | |
| from ..net.geometry import orthogonal | |
| from pytorch3d.renderer.mesh import rasterize_meshes | |
| from .render_utils import Pytorch3dRasterizer | |
| from pytorch3d.structures import Meshes | |
| import cv2 | |
| from PIL import Image | |
| from tqdm import tqdm | |
| import os | |
| from termcolor import colored | |
| 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_eval(opt, net, cuda, data, resolution=None): | |
| resolution = opt.resolution if resolution is None else resolution | |
| 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: | |
| verts, faces, _, _ = reconstruction_faster(net, | |
| cuda, | |
| calib_tensor, | |
| resolution, | |
| b_min, | |
| b_max, | |
| use_octree=False) | |
| except Exception as e: | |
| print(e) | |
| print('Can not create marching cubes at this time.') | |
| verts, faces = None, None | |
| return verts, faces | |
| def gen_mesh(opt, net, cuda, data, save_path, resolution=None): | |
| resolution = opt.resolution if resolution is None else resolution | |
| 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_faster(net, cuda, calib_tensor, | |
| resolution, b_min, b_max) | |
| 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 = netG.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.') | |
| verts, faces, color = None, None, None | |
| return verts, faces, color | |
| 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_faster(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.') | |
| verts, faces, color = None, None, None | |
| return verts, faces, color | |
| 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_metrics(opt, net, cuda, dataset, num_tests, | |
| # resolution=128, sampled_points=1000, use_kaolin=True): | |
| # if num_tests > len(dataset): | |
| # num_tests = len(dataset) | |
| # with torch.no_grad(): | |
| # chamfer_arr, p2s_arr = [], [] | |
| # for idx in tqdm(range(num_tests)): | |
| # data = dataset[idx * len(dataset) // num_tests] | |
| # verts, faces = gen_mesh_eval(opt, net, cuda, data, resolution) | |
| # if verts is None: | |
| # continue | |
| # mesh_gt = trimesh.load(data['mesh_path']) | |
| # mesh_gt = mesh_gt.split(only_watertight=False) | |
| # comp_num = [mesh.vertices.shape[0] for mesh in mesh_gt] | |
| # mesh_gt = mesh_gt[comp_num.index(max(comp_num))] | |
| # mesh_pred = trimesh.Trimesh(verts, faces) | |
| # gt_surface_pts, _ = trimesh.sample.sample_surface_even( | |
| # mesh_gt, sampled_points) | |
| # pred_surface_pts, _ = trimesh.sample.sample_surface_even( | |
| # mesh_pred, sampled_points) | |
| # if use_kaolin and has_kaolin: | |
| # kal_mesh_gt = kal.rep.TriangleMesh.from_tensors( | |
| # torch.tensor(mesh_gt.vertices).float().to(device=cuda), | |
| # torch.tensor(mesh_gt.faces).long().to(device=cuda)) | |
| # kal_mesh_pred = kal.rep.TriangleMesh.from_tensors( | |
| # torch.tensor(mesh_pred.vertices).float().to(device=cuda), | |
| # torch.tensor(mesh_pred.faces).long().to(device=cuda)) | |
| # kal_distance_0 = kal.metrics.mesh.point_to_surface( | |
| # torch.tensor(pred_surface_pts).float().to(device=cuda), kal_mesh_gt) | |
| # kal_distance_1 = kal.metrics.mesh.point_to_surface( | |
| # torch.tensor(gt_surface_pts).float().to(device=cuda), kal_mesh_pred) | |
| # dist_gt_pred = torch.sqrt(kal_distance_0).cpu().numpy() | |
| # dist_pred_gt = torch.sqrt(kal_distance_1).cpu().numpy() | |
| # else: | |
| # try: | |
| # _, dist_pred_gt, _ = trimesh.proximity.closest_point(mesh_pred, gt_surface_pts) | |
| # _, dist_gt_pred, _ = trimesh.proximity.closest_point(mesh_gt, pred_surface_pts) | |
| # except Exception as e: | |
| # print (e) | |
| # continue | |
| # chamfer_dist = 0.5 * (dist_pred_gt.mean() + dist_gt_pred.mean()) | |
| # p2s_dist = dist_pred_gt.mean() | |
| # chamfer_arr.append(chamfer_dist) | |
| # p2s_arr.append(p2s_dist) | |
| # return np.average(chamfer_arr), np.average(p2s_arr) | |
| 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) | |
| # pytorch lightning training related fucntions | |
| def query_func(opt, netG, features, points, proj_matrix=None): | |
| ''' | |
| - points: size of (bz, N, 3) | |
| - proj_matrix: size of (bz, 4, 4) | |
| return: size of (bz, 1, N) | |
| ''' | |
| assert len(points) == 1 | |
| samples = points.repeat(opt.num_views, 1, 1) | |
| samples = samples.permute(0, 2, 1) # [bz, 3, N] | |
| # view specific query | |
| if proj_matrix is not None: | |
| samples = orthogonal(samples, proj_matrix) | |
| calib_tensor = torch.stack([torch.eye(4).float()], dim=0).type_as(samples) | |
| preds = netG.query(features=features, | |
| points=samples, | |
| calibs=calib_tensor, | |
| regressor=netG.if_regressor) | |
| if type(preds) is list: | |
| preds = preds[0] | |
| return preds | |
| def isin(ar1, ar2): | |
| return (ar1[..., None] == ar2).any(-1) | |
| def in1d(ar1, ar2): | |
| mask = ar2.new_zeros((max(ar1.max(), ar2.max()) + 1, ), dtype=torch.bool) | |
| mask[ar2.unique()] = True | |
| return mask[ar1] | |
| def get_visibility(xy, z, faces): | |
| """get the visibility of vertices | |
| Args: | |
| xy (torch.tensor): [N,2] | |
| z (torch.tensor): [N,1] | |
| faces (torch.tensor): [N,3] | |
| size (int): resolution of rendered image | |
| """ | |
| xyz = torch.cat((xy, -z), dim=1) | |
| xyz = (xyz + 1.0) / 2.0 | |
| faces = faces.long() | |
| rasterizer = Pytorch3dRasterizer(image_size=2**12) | |
| meshes_screen = Meshes(verts=xyz[None, ...], faces=faces[None, ...]) | |
| raster_settings = rasterizer.raster_settings | |
| pix_to_face, zbuf, bary_coords, dists = rasterize_meshes( | |
| meshes_screen, | |
| image_size=raster_settings.image_size, | |
| blur_radius=raster_settings.blur_radius, | |
| faces_per_pixel=raster_settings.faces_per_pixel, | |
| bin_size=raster_settings.bin_size, | |
| max_faces_per_bin=raster_settings.max_faces_per_bin, | |
| perspective_correct=raster_settings.perspective_correct, | |
| cull_backfaces=raster_settings.cull_backfaces, | |
| ) | |
| vis_vertices_id = torch.unique(faces[torch.unique(pix_to_face), :]) | |
| vis_mask = torch.zeros(size=(z.shape[0], 1)) | |
| vis_mask[vis_vertices_id] = 1.0 | |
| # print("------------------------\n") | |
| # print(f"keep points : {vis_mask.sum()/len(vis_mask)}") | |
| return vis_mask | |
| def batch_mean(res, key): | |
| # recursive mean for multilevel dicts | |
| return torch.stack([ | |
| x[key] if isinstance(x, dict) else batch_mean(x, key) for x in res | |
| ]).mean() | |
| def tf_log_convert(log_dict): | |
| new_log_dict = log_dict.copy() | |
| for k, v in log_dict.items(): | |
| new_log_dict[k.replace("_", "/")] = v | |
| del new_log_dict[k] | |
| return new_log_dict | |
| def bar_log_convert(log_dict, name=None, rot=None): | |
| from decimal import Decimal | |
| new_log_dict = {} | |
| if name is not None: | |
| new_log_dict['name'] = name[0] | |
| if rot is not None: | |
| new_log_dict['rot'] = rot[0] | |
| for k, v in log_dict.items(): | |
| color = "yellow" | |
| if 'loss' in k: | |
| color = "red" | |
| k = k.replace("loss", "L") | |
| elif 'acc' in k: | |
| color = "green" | |
| k = k.replace("acc", "A") | |
| elif 'iou' in k: | |
| color = "green" | |
| k = k.replace("iou", "I") | |
| elif 'prec' in k: | |
| color = "green" | |
| k = k.replace("prec", "P") | |
| elif 'recall' in k: | |
| color = "green" | |
| k = k.replace("recall", "R") | |
| if 'lr' not in k: | |
| new_log_dict[colored(k.split("_")[1], | |
| color)] = colored(f"{v:.3f}", color) | |
| else: | |
| new_log_dict[colored(k.split("_")[1], | |
| color)] = colored(f"{Decimal(str(v)):.1E}", | |
| color) | |
| if 'loss' in new_log_dict.keys(): | |
| del new_log_dict['loss'] | |
| return new_log_dict | |
| def accumulate(outputs, rot_num, split): | |
| hparam_log_dict = {} | |
| metrics = outputs[0].keys() | |
| datasets = split.keys() | |
| for dataset in datasets: | |
| for metric in metrics: | |
| keyword = f"hparam/{dataset}-{metric}" | |
| if keyword not in hparam_log_dict.keys(): | |
| hparam_log_dict[keyword] = 0 | |
| for idx in range(split[dataset][0] * rot_num, | |
| split[dataset][1] * rot_num): | |
| hparam_log_dict[keyword] += outputs[idx][metric] | |
| hparam_log_dict[keyword] /= (split[dataset][1] - | |
| split[dataset][0]) * rot_num | |
| print(colored(hparam_log_dict, "green")) | |
| return hparam_log_dict | |
| def calc_error_N(outputs, targets): | |
| """calculate the error of normal (IGR) | |
| Args: | |
| outputs (torch.tensor): [B, 3, N] | |
| target (torch.tensor): [B, N, 3] | |
| # manifold loss and grad_loss in IGR paper | |
| grad_loss = ((nonmnfld_grad.norm(2, dim=-1) - 1) ** 2).mean() | |
| normals_loss = ((mnfld_grad - normals).abs()).norm(2, dim=1).mean() | |
| Returns: | |
| torch.tensor: error of valid normals on the surface | |
| """ | |
| # outputs = torch.tanh(-outputs.permute(0,2,1).reshape(-1,3)) | |
| outputs = -outputs.permute(0, 2, 1).reshape(-1, 1) | |
| targets = targets.reshape(-1, 3)[:, 2:3] | |
| with_normals = targets.sum(dim=1).abs() > 0.0 | |
| # eikonal loss | |
| grad_loss = ((outputs[with_normals].norm(2, dim=-1) - 1)**2).mean() | |
| # normals loss | |
| normal_loss = (outputs - targets)[with_normals].abs().norm(2, dim=1).mean() | |
| return grad_loss * 0.0 + normal_loss | |
| def calc_knn_acc(preds, carn_verts, labels, pick_num): | |
| """calculate knn accuracy | |
| Args: | |
| preds (torch.tensor): [B, 3, N] | |
| carn_verts (torch.tensor): [SMPLX_V_num, 3] | |
| labels (torch.tensor): [B, N_knn, N] | |
| """ | |
| N_knn_full = labels.shape[1] | |
| preds = preds.permute(0, 2, 1).reshape(-1, 3) | |
| labels = labels.permute(0, 2, 1).reshape(-1, N_knn_full) # [BxN, num_knn] | |
| labels = labels[:, :pick_num] | |
| dist = torch.cdist(preds, carn_verts, p=2) # [BxN, SMPL_V_num] | |
| knn = dist.topk(k=pick_num, dim=1, largest=False)[1] # [BxN, num_knn] | |
| cat_mat = torch.sort(torch.cat((knn, labels), dim=1))[0] | |
| bool_col = torch.zeros_like(cat_mat)[:, 0] | |
| for i in range(pick_num * 2 - 1): | |
| bool_col += cat_mat[:, i] == cat_mat[:, i + 1] | |
| acc = (bool_col > 0).sum() / len(bool_col) | |
| return acc | |
| def calc_acc_seg(output, target, num_multiseg): | |
| from pytorch_lightning.metrics import Accuracy | |
| return Accuracy()(output.reshape(-1, num_multiseg).cpu(), | |
| target.flatten().cpu()) | |
| def add_watermark(imgs, titles): | |
| # Write some Text | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| bottomLeftCornerOfText = (350, 50) | |
| bottomRightCornerOfText = (800, 50) | |
| fontScale = 1 | |
| fontColor = (1.0, 1.0, 1.0) | |
| lineType = 2 | |
| for i in range(len(imgs)): | |
| title = titles[i + 1] | |
| cv2.putText(imgs[i], title, bottomLeftCornerOfText, font, fontScale, | |
| fontColor, lineType) | |
| if i == 0: | |
| cv2.putText(imgs[i], str(titles[i][0]), bottomRightCornerOfText, | |
| font, fontScale, fontColor, lineType) | |
| result = np.concatenate(imgs, axis=0).transpose(2, 0, 1) | |
| return result | |
| def make_test_gif(img_dir): | |
| if img_dir is not None and len(os.listdir(img_dir)) > 0: | |
| for dataset in os.listdir(img_dir): | |
| for subject in sorted(os.listdir(osp.join(img_dir, dataset))): | |
| img_lst = [] | |
| im1 = None | |
| for file in sorted( | |
| os.listdir(osp.join(img_dir, dataset, subject))): | |
| if file[-3:] not in ['obj', 'gif']: | |
| img_path = os.path.join(img_dir, dataset, subject, | |
| file) | |
| if im1 == None: | |
| im1 = Image.open(img_path) | |
| else: | |
| img_lst.append(Image.open(img_path)) | |
| print(os.path.join(img_dir, dataset, subject, "out.gif")) | |
| im1.save(os.path.join(img_dir, dataset, subject, "out.gif"), | |
| save_all=True, | |
| append_images=img_lst, | |
| duration=500, | |
| loop=0) | |
| def export_cfg(logger, cfg): | |
| cfg_export_file = osp.join(logger.save_dir, logger.name, | |
| f"version_{logger.version}", "cfg.yaml") | |
| if not osp.exists(cfg_export_file): | |
| os.makedirs(osp.dirname(cfg_export_file), exist_ok=True) | |
| with open(cfg_export_file, "w+") as file: | |
| _ = yaml.dump(cfg, file) | |