# -*- 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)