nerf_plus_plus/utils.py
2020-10-11 21:33:31 -04:00

206 lines
5.6 KiB
Python

import torch
# import torch.nn as nn
# import torch.nn.functional as F
import numpy as np
HUGE_NUMBER = 1e10
TINY_NUMBER = 1e-6 # float32 only has 7 decimal digits precision
# Misc utils
# work on tensors
# img2mse = lambda x, y: torch.mean((x - y) * (x - y))
def img2mse(x, y, mask=None):
if mask is None:
return torch.mean((x - y) * (x - y))
else:
return torch.sum((x - y) * (x - y) * mask.unsqueeze(-1)) / (torch.sum(mask) * x.shape[-1] + TINY_NUMBER)
img_HWC2CHW = lambda x: x.permute(2, 0, 1)
gray2rgb = lambda x: x.unsqueeze(2).repeat(1, 1, 3)
def normalize(x):
min = x.min()
max = x.max()
return (x - min) / ((max - min) + TINY_NUMBER)
to8b = lambda x: (255 * np.clip(x, 0, 1)).astype(np.uint8)
# gray2rgb = lambda x: np.tile(x[:,:,np.newaxis], (1, 1, 3))
mse2psnr = lambda x: -10. * np.log(x+TINY_NUMBER) / np.log(10.)
#
# def normalize(x):
# x_min, x_max = np.percentile(x, (0.5, 99.5))
# x = np.clip(x, x_min, x_max)
# x = (x - x_min) / (x_max - x_min)
# return x
########################################################################################################################
#
########################################################################################################################
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure
import matplotlib as mpl
from matplotlib import cm
import cv2
def get_vertical_colorbar(h, vmin, vmax, cmap_name='jet', label=None):
'''
:param w: pixels
:param h: pixels
:param vmin: min value
:param vmax: max value
:param cmap_name:
:param label
:return:
'''
fig = Figure(figsize=(1.2, 8), dpi=100)
fig.subplots_adjust(right=1.5)
canvas = FigureCanvasAgg(fig)
# Do some plotting.
ax = fig.add_subplot(111)
cmap = cm.get_cmap(cmap_name)
norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
tick_cnt = 6
tick_loc = np.linspace(vmin, vmax, tick_cnt)
cb1 = mpl.colorbar.ColorbarBase(ax, cmap=cmap,
norm=norm,
ticks=tick_loc,
orientation='vertical')
tick_label = ['{:3.2f}'.format(x) for x in tick_loc]
cb1.set_ticklabels(tick_label)
cb1.ax.tick_params(labelsize=18, rotation=0)
if label is not None:
cb1.set_label(label)
fig.tight_layout()
# # debug
# fig.savefig("debug3.png")
canvas.draw()
s, (width, height) = canvas.print_to_buffer()
im = np.frombuffer(s, np.uint8).reshape((height, width, 4))
im = im[:, :, :3].astype(np.float32) / 255.
if h != im.shape[0]:
w = int(im.shape[1] / im.shape[0] * h)
im = cv2.resize(im, (w, h), interpolation=cv2.INTER_AREA)
return im
# def colorize_np(x, cmap_name='jet', append_cbar=False):
# vmin = x.min()
# vmax = x.max() + TINY_NUMBER
# x = (x - vmin) / (vmax - vmin)
# # x = np.clip(x, 0., 1.)
# cmap = cm.get_cmap(cmap_name)
# x_new = cmap(x)[:, :, :3]
# cbar = get_vertical_colorbar(h=x.shape[0], vmin=vmin, vmax=vmax, cmap_name=cmap_name)
# if append_cbar:
# x_new = np.concatenate((x_new, np.zeros_like(x_new[:, :5, :]), cbar), axis=1)
# return x_new
# else:
# return x_new, cbar
# # tensor
# def colorize(x, cmap_name='jet', append_cbar=False):
# x = x.numpy()
# x, cbar = colorize_np(x, cmap_name)
# if append_cbar:
# x = np.concatenate((x, np.zeros_like(x[:, :5, :]), cbar), axis=1)
# x = torch.from_numpy(x)
# return x
def colorize_np(x, cmap_name='jet', mask=None, append_cbar=False):
if mask is not None:
# vmin, vmax = np.percentile(x[mask], (1, 99))
vmin = np.min(x[mask])
vmax = np.max(x[mask])
vmin = vmin - np.abs(vmin) * 0.01
x[np.logical_not(mask)] = vmin
x = np.clip(x, vmin, vmax)
# print(vmin, vmax)
else:
vmin = x.min()
vmax = x.max() + TINY_NUMBER
x = (x - vmin) / (vmax - vmin)
# x = np.clip(x, 0., 1.)
cmap = cm.get_cmap(cmap_name)
x_new = cmap(x)[:, :, :3]
if mask is not None:
mask = np.float32(mask[:, :, np.newaxis])
x_new = x_new * mask + np.zeros_like(x_new) * (1. - mask)
cbar = get_vertical_colorbar(h=x.shape[0], vmin=vmin, vmax=vmax, cmap_name=cmap_name)
if append_cbar:
x_new = np.concatenate((x_new, np.zeros_like(x_new[:, :5, :]), cbar), axis=1)
return x_new
else:
return x_new, cbar
# tensor
def colorize(x, cmap_name='jet', append_cbar=False, mask=None):
x = x.numpy()
if mask is not None:
mask = mask.numpy().astype(dtype=np.bool)
x, cbar = colorize_np(x, cmap_name, mask)
if append_cbar:
x = np.concatenate((x, np.zeros_like(x[:, :5, :]), cbar), axis=1)
x = torch.from_numpy(x)
return x
if __name__ == '__main__':
# # cbar = get_vertical_colorbar(h=512, vmin=0.1, vmax=5, cmap_name='jet')
# # cbar = cbar[:, :, :3]
# import imageio
#
# # imageio.imwrite('./debug.png', cbar)
#
# x = torch.rand(512, 512)
# x = colorize(x, append_cbar=True)
#
# x = np.uint8(x.numpy() * 255.)
#
# import imageio
# imageio.imwrite('./debug.png', x)
import os
import imageio
img_dir = '/phoenix/S7/kz298/latest_work/nerf/logs/dtu_scan9_3_nearfar/renderonly_train_200001'
all_imgs = []
for item in sorted(os.listdir(img_dir)):
if item[-4:] == '.png':
fpath = os.path.join(img_dir, item)
all_imgs.append(imageio.imread(fpath))
imageio.mimwrite(os.path.join(img_dir, 'video.mp4'), all_imgs, fps=3, quality=8)