Spaces:

qitaoz
/

SparseAGS

Paused

init commit

4f54ccd verified almost 2 years ago

9.02 kB

import os
import re
import cv2
import csv
import json
import math
import tqdm
import shutil
import argparse
import numpy as np
from PIL import Image
from omegaconf import OmegaConf
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F
import nvdiffrast.torch as dr

from kiui.mesh import Mesh
from kiui.cam import OrbitCamera
from kiui.op import safe_normalize
from kiui.lpips import LPIPS

import sys 
from sparseags.mesh_utils.mesh_renderer import Renderer
from sparseags.cam_utils import orbit_camera, OrbitCamera
from sparseags.render_utils.gs_renderer import CustomCamera


class GUI:
	def __init__(self, opt):
		self.opt = opt
		self.W = opt.W
		self.H = opt.H
		self.wogui = opt.wogui # disable gui and run in cmd
		self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
		self.bg_color = torch.ones(3, dtype=torch.float32).cuda() # default white bg
		# self.bg_color = torch.zeros(3, dtype=torch.float32).cuda() # black bg

		self.render_buffer = np.zeros((self.W, self.H, 3), dtype=np.float32)
		self.need_update = True # camera moved, should reset accumulation
		self.light_dir = np.array([0, 0])
		self.ambient_ratio = 0.5

		# auto-rotate
		self.auto_rotate_cam = False
		self.auto_rotate_light = False
		
		self.mode = opt.mode
		self.render_modes = ['albedo', 'depth', 'normal', 'lambertian']

		# load mesh
		self.mesh = Mesh.load(opt.mesh, front_dir=opt.front_dir)

		if not opt.force_cuda_rast and (self.wogui or os.name == 'nt'):
			self.glctx = dr.RasterizeGLContext()
		else:
			self.glctx = dr.RasterizeCudaContext()
	
	def step(self):

		if not self.need_update:
			return
	
		starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
		starter.record()

		# do MVP for vertices
		pose = torch.from_numpy(self.cam.pose.astype(np.float32)).cuda()
		proj = torch.from_numpy(self.cam.perspective.astype(np.float32)).cuda()
		
		v_cam = torch.matmul(F.pad(self.mesh.v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
		v_clip = v_cam @ proj.T

		rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (self.H, self.W))

		alpha = (rast[..., 3:] > 0).float()
		alpha = dr.antialias(alpha, rast, v_clip, self.mesh.f).squeeze(0).clamp(0, 1) # [H, W, 3]
		
		if self.mode == 'depth':
			depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
			depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-20)
			buffer = depth.squeeze(0).detach().cpu().numpy().repeat(3, -1) # [H, W, 3]
		else:
			# use vertex color if exists
			if self.mesh.vc is not None:
				albedo, _ = dr.interpolate(self.mesh.vc.unsqueeze(0).contiguous(), rast, self.mesh.f)
			# use texture image
			else:
				texc, _ = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft)
				albedo = dr.texture(self.mesh.albedo.unsqueeze(0), texc, filter_mode='linear') # [1, H, W, 3]

			albedo = torch.where(rast[..., 3:] > 0, albedo, torch.tensor(0).to(albedo.device)) # remove background
			albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).clamp(0, 1) # [1, H, W, 3]
			if self.mode == 'albedo':
				albedo = albedo * alpha + self.bg_color * (1 - alpha)
				buffer = albedo[0].detach().cpu().numpy()
			else:
				normal, _ = dr.interpolate(self.mesh.vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
				normal = safe_normalize(normal)
				if self.mode == 'normal':
					normal_image = (normal[0] + 1) / 2
					normal_image = torch.where(rast[..., 3:] > 0, normal_image, torch.tensor(1).to(normal_image.device)) # remove background
					buffer = normal_image.detach().cpu().numpy()
				elif self.mode == 'lambertian':
					light_d = np.deg2rad(self.light_dir)
					light_d = np.array([
						np.cos(light_d[0]) * np.sin(light_d[1]),
						-np.sin(light_d[0]),
						np.cos(light_d[0]) * np.cos(light_d[1]),
					], dtype=np.float32)
					light_d = torch.from_numpy(light_d).to(albedo.device)
					lambertian = self.ambient_ratio + (1 - self.ambient_ratio)  * (normal @ light_d).float().clamp(min=0)
					albedo = (albedo * lambertian.unsqueeze(-1)) * alpha + self.bg_color * (1 - alpha)
					buffer = albedo[0].detach().cpu().numpy()

		ender.record()
		torch.cuda.synchronize()
		t = starter.elapsed_time(ender)

		self.render_buffer = buffer
		self.need_update = False

		if self.auto_rotate_cam:
			self.cam.orbit(5, 0)
			self.need_update = True
		
		if self.auto_rotate_light:
			self.light_dir[1] += 3
			self.need_update = True


def vis_output(camera_data, mesh_path=None, save_path=None, num_views=8):
	parser = argparse.ArgumentParser()
	parser.add_argument('--front_dir', type=str, default='+z', help="mesh front-facing dir")
	parser.add_argument('--mode', default='albedo', type=str, choices=['lambertian', 'albedo', 'normal', 'depth'], help="rendering mode")
	parser.add_argument('--W', type=int, default=256, help="GUI width")
	parser.add_argument('--H', type=int, default=256, help="GUI height")
	parser.add_argument('--radius', type=float, default=3, help="default GUI camera radius from center")
	parser.add_argument('--fovy', type=float, default=49.1, help="default GUI camera fovy")
	parser.add_argument("--wogui", type=bool, default=True, help="disable all dpg GUI")
	parser.add_argument("--force_cuda_rast", action='store_true', help="force to use RasterizeCudaContext.")
	parser.add_argument('--elevation', type=int, default=0, help="rendering elevation")
	parser.add_argument('--save_video', type=str, default=None, help="path to save rendered video")
	parser.add_argument('--idx', type=int, default=0, help="GUI height")
	parser.add_argument('--config', default='configs/navi.yaml', type=str, help='Path to config directory, which contains image.yaml')
	args, extras = parser.parse_known_args()

	# override default config from cli
	opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))
	data = camera_data

	cameras = [CustomCamera(cam_params) for cam_params in data.values()]
	cams = [(cam.c2w, cam.perspective, cam.focal_length) for cam in cameras]
	img_paths = [v["filepath"] for k, v in data.items()]

	opt.mesh = mesh_path
	opt.trainable_texture = False
	renderer = Renderer(opt).to(torch.device("cuda"))

	lpips_loss = LPIPS(net='vgg').cuda()
	mse_losses = []
	lpips_losses = []
	flags = [int(v["flag"]) for k, v in data.items()]
	images = np.zeros((2, num_views, 256, 256, 3), dtype=np.uint8)

	for i in tqdm.tqdm(range(len(cams))):

		img_path = img_paths[i]

		img = Image.open(img_path)
		if img.mode == 'RGBA':
			img = np.asarray(img, dtype=np.uint8).copy()
			img[img[:, :, -1] <= 255*0.9] = [255., 255., 255., 255.] # thresholding background
			img = img[:, :, :3]

		gt_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0).cuda() / 255.0

		images[0, i] = img

		with torch.no_grad():
			out = renderer.render(*cams[i][:2], 256, 256, ssaa=1)

		# rgb loss
		image = (out["image"].detach().cpu().numpy() * 255).astype(np.uint8)
		pred_tensor = out["image"].permute(2, 0, 1).float().unsqueeze(0).cuda()
		# obj_scale = ((out["alpha"] > 0) & (out["viewcos"] > 0.5)).detach().sum().float()
		obj_scale = (out["alpha"] > 0).detach().sum().float()
		obj_scale /= 256 ** 2
		
		images[1, i] = image
		with torch.no_grad():
			mse_losses.append(F.mse_loss(pred_tensor, gt_tensor).squeeze().cpu().numpy() / obj_scale.item())
			lpips_losses.append(lpips_loss(pred_tensor, gt_tensor).squeeze().cpu().numpy() / obj_scale.item())

	mean_mse = np.mean(np.array(mse_losses)[:num_views])
	mean_lpips = np.mean(np.array(lpips_losses)[:num_views])

	num_frames = 2 * num_views
	cmap = plt.get_cmap("hsv")
	num_rows = 2
	num_cols = num_views
	plt.subplots_adjust(top=0.2)
	figsize = (num_cols * 2, num_rows * 2.2)
	fig, axs = plt.subplots(num_rows, num_cols, figsize=figsize)
	fig.suptitle(f"Avg MSE: {mean_mse:.4f}, Avg LPIPS: {mean_lpips:.4f}", fontsize=16, y=0.97)
	axs = axs.flatten()
	for i in range(num_rows * num_cols):
		if i < num_frames:
			axs[i].imshow(images.reshape(-1, 256, 256, 3)[i])
			for s in ["bottom", "top", "left", "right"]:
				if i % num_views <= num_views - 1:
					if not flags[i%num_views]:
						axs[i].spines[s].set_color("red")
					else:
						axs[i].spines[s].set_color("green")
				else:
					axs[i].spines[s].set_color(cmap(i / (num_frames)))
				axs[i].spines[s].set_linewidth(5)
			axs[i].set_xticks([])
			axs[i].set_yticks([])

			if i >= num_views:
				axs[i].set_xlabel(f'MSE: {mse_losses[i%num_views]:.4f}\nLPIPS: {lpips_losses[i%num_views]:.4f}', fontsize=10)
		else:
			axs[i].axis("off")
	plt.tight_layout()
	plt.savefig(save_path)
	plt.close(fig)
	print(f"Visualization file written to {save_path}")

	out_dir = save_path.replace('vis.png', 'reprojections') 
	os.makedirs(out_dir, exist_ok=True)

	for i in range(num_views):
		gt = Image.fromarray(images[0, i])
		pred = Image.fromarray(images[1, i])
		gt.save(os.path.join(out_dir, f"gt_{i}.png"))
		pred.save(os.path.join(out_dir, f"pred_{i}.png"))

	return np.array(lpips_losses), np.array(mse_losses)

	

	import os
	import re
	import cv2
	import csv
	import json
	import math
	import tqdm
	import shutil
	import argparse
	import numpy as np
	from PIL import Image
	from omegaconf import OmegaConf
	import matplotlib.pyplot as plt

	import torch
	import torch.nn.functional as F
	import nvdiffrast.torch as dr

	from kiui.mesh import Mesh
	from kiui.cam import OrbitCamera
	from kiui.op import safe_normalize
	from kiui.lpips import LPIPS

	import sys
	from sparseags.mesh_utils.mesh_renderer import Renderer
	from sparseags.cam_utils import orbit_camera, OrbitCamera
	from sparseags.render_utils.gs_renderer import CustomCamera


	class GUI:
	def __init__(self, opt):
	self.opt = opt
	self.W = opt.W
	self.H = opt.H
	self.wogui = opt.wogui # disable gui and run in cmd
	self.cam = OrbitCamera(opt.W, opt.H, r=opt.radius, fovy=opt.fovy)
	self.bg_color = torch.ones(3, dtype=torch.float32).cuda() # default white bg
	# self.bg_color = torch.zeros(3, dtype=torch.float32).cuda() # black bg

	self.render_buffer = np.zeros((self.W, self.H, 3), dtype=np.float32)
	self.need_update = True # camera moved, should reset accumulation
	self.light_dir = np.array([0, 0])
	self.ambient_ratio = 0.5

	# auto-rotate
	self.auto_rotate_cam = False
	self.auto_rotate_light = False

	self.mode = opt.mode
	self.render_modes = ['albedo', 'depth', 'normal', 'lambertian']

	# load mesh
	self.mesh = Mesh.load(opt.mesh, front_dir=opt.front_dir)

	if not opt.force_cuda_rast and (self.wogui or os.name == 'nt'):
	self.glctx = dr.RasterizeGLContext()
	else:
	self.glctx = dr.RasterizeCudaContext()

	def step(self):

	if not self.need_update:
	return

	starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
	starter.record()

	# do MVP for vertices
	pose = torch.from_numpy(self.cam.pose.astype(np.float32)).cuda()
	proj = torch.from_numpy(self.cam.perspective.astype(np.float32)).cuda()

	v_cam = torch.matmul(F.pad(self.mesh.v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
	v_clip = v_cam @ proj.T

	rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (self.H, self.W))

	alpha = (rast[..., 3:] > 0).float()
	alpha = dr.antialias(alpha, rast, v_clip, self.mesh.f).squeeze(0).clamp(0, 1) # [H, W, 3]

	if self.mode == 'depth':
	depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
	depth = (depth - depth.min()) / (depth.max() - depth.min() + 1e-20)
	buffer = depth.squeeze(0).detach().cpu().numpy().repeat(3, -1) # [H, W, 3]
	else:
	# use vertex color if exists
	if self.mesh.vc is not None:
	albedo, _ = dr.interpolate(self.mesh.vc.unsqueeze(0).contiguous(), rast, self.mesh.f)
	# use texture image
	else:
	texc, _ = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft)
	albedo = dr.texture(self.mesh.albedo.unsqueeze(0), texc, filter_mode='linear') # [1, H, W, 3]

	albedo = torch.where(rast[..., 3:] > 0, albedo, torch.tensor(0).to(albedo.device)) # remove background
	albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).clamp(0, 1) # [1, H, W, 3]
	if self.mode == 'albedo':
	albedo = albedo * alpha + self.bg_color * (1 - alpha)
	buffer = albedo[0].detach().cpu().numpy()
	else:
	normal, _ = dr.interpolate(self.mesh.vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
	normal = safe_normalize(normal)
	if self.mode == 'normal':
	normal_image = (normal[0] + 1) / 2
	normal_image = torch.where(rast[..., 3:] > 0, normal_image, torch.tensor(1).to(normal_image.device)) # remove background
	buffer = normal_image.detach().cpu().numpy()
	elif self.mode == 'lambertian':
	light_d = np.deg2rad(self.light_dir)
	light_d = np.array([
	np.cos(light_d[0]) * np.sin(light_d[1]),
	-np.sin(light_d[0]),
	np.cos(light_d[0]) * np.cos(light_d[1]),
	], dtype=np.float32)
	light_d = torch.from_numpy(light_d).to(albedo.device)
	lambertian = self.ambient_ratio + (1 - self.ambient_ratio) * (normal @ light_d).float().clamp(min=0)
	albedo = (albedo * lambertian.unsqueeze(-1)) * alpha + self.bg_color * (1 - alpha)
	buffer = albedo[0].detach().cpu().numpy()

	ender.record()
	torch.cuda.synchronize()
	t = starter.elapsed_time(ender)

	self.render_buffer = buffer
	self.need_update = False

	if self.auto_rotate_cam:
	self.cam.orbit(5, 0)
	self.need_update = True

	if self.auto_rotate_light:
	self.light_dir[1] += 3
	self.need_update = True


	def vis_output(camera_data, mesh_path=None, save_path=None, num_views=8):
	parser = argparse.ArgumentParser()
	parser.add_argument('--front_dir', type=str, default='+z', help="mesh front-facing dir")
	parser.add_argument('--mode', default='albedo', type=str, choices=['lambertian', 'albedo', 'normal', 'depth'], help="rendering mode")
	parser.add_argument('--W', type=int, default=256, help="GUI width")
	parser.add_argument('--H', type=int, default=256, help="GUI height")
	parser.add_argument('--radius', type=float, default=3, help="default GUI camera radius from center")
	parser.add_argument('--fovy', type=float, default=49.1, help="default GUI camera fovy")
	parser.add_argument("--wogui", type=bool, default=True, help="disable all dpg GUI")
	parser.add_argument("--force_cuda_rast", action='store_true', help="force to use RasterizeCudaContext.")
	parser.add_argument('--elevation', type=int, default=0, help="rendering elevation")
	parser.add_argument('--save_video', type=str, default=None, help="path to save rendered video")
	parser.add_argument('--idx', type=int, default=0, help="GUI height")
	parser.add_argument('--config', default='configs/navi.yaml', type=str, help='Path to config directory, which contains image.yaml')
	args, extras = parser.parse_known_args()

	# override default config from cli
	opt = OmegaConf.merge(OmegaConf.load(args.config), OmegaConf.from_cli(extras))
	data = camera_data

	cameras = [CustomCamera(cam_params) for cam_params in data.values()]
	cams = [(cam.c2w, cam.perspective, cam.focal_length) for cam in cameras]
	img_paths = [v["filepath"] for k, v in data.items()]

	opt.mesh = mesh_path
	opt.trainable_texture = False
	renderer = Renderer(opt).to(torch.device("cuda"))

	lpips_loss = LPIPS(net='vgg').cuda()
	mse_losses = []
	lpips_losses = []
	flags = [int(v["flag"]) for k, v in data.items()]
	images = np.zeros((2, num_views, 256, 256, 3), dtype=np.uint8)

	for i in tqdm.tqdm(range(len(cams))):

	img_path = img_paths[i]

	img = Image.open(img_path)
	if img.mode == 'RGBA':
	img = np.asarray(img, dtype=np.uint8).copy()
	img[img[:, :, -1] <= 255*0.9] = [255., 255., 255., 255.] # thresholding background
	img = img[:, :, :3]

	gt_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0).cuda() / 255.0

	images[0, i] = img

	with torch.no_grad():
	out = renderer.render(*cams[i][:2], 256, 256, ssaa=1)

	# rgb loss
	image = (out["image"].detach().cpu().numpy() * 255).astype(np.uint8)
	pred_tensor = out["image"].permute(2, 0, 1).float().unsqueeze(0).cuda()
	# obj_scale = ((out["alpha"] > 0) & (out["viewcos"] > 0.5)).detach().sum().float()
	obj_scale = (out["alpha"] > 0).detach().sum().float()
	obj_scale /= 256 ** 2

	images[1, i] = image
	with torch.no_grad():
	mse_losses.append(F.mse_loss(pred_tensor, gt_tensor).squeeze().cpu().numpy() / obj_scale.item())
	lpips_losses.append(lpips_loss(pred_tensor, gt_tensor).squeeze().cpu().numpy() / obj_scale.item())

	mean_mse = np.mean(np.array(mse_losses)[:num_views])
	mean_lpips = np.mean(np.array(lpips_losses)[:num_views])

	num_frames = 2 * num_views
	cmap = plt.get_cmap("hsv")
	num_rows = 2
	num_cols = num_views
	plt.subplots_adjust(top=0.2)
	figsize = (num_cols * 2, num_rows * 2.2)
	fig, axs = plt.subplots(num_rows, num_cols, figsize=figsize)
	fig.suptitle(f"Avg MSE: {mean_mse:.4f}, Avg LPIPS: {mean_lpips:.4f}", fontsize=16, y=0.97)
	axs = axs.flatten()
	for i in range(num_rows * num_cols):
	if i < num_frames:
	axs[i].imshow(images.reshape(-1, 256, 256, 3)[i])
	for s in ["bottom", "top", "left", "right"]:
	if i % num_views <= num_views - 1:
	if not flags[i%num_views]:
	axs[i].spines[s].set_color("red")
	else:
	axs[i].spines[s].set_color("green")
	else:
	axs[i].spines[s].set_color(cmap(i / (num_frames)))
	axs[i].spines[s].set_linewidth(5)
	axs[i].set_xticks([])
	axs[i].set_yticks([])

	if i >= num_views:
	axs[i].set_xlabel(f'MSE: {mse_losses[i%num_views]:.4f}\nLPIPS: {lpips_losses[i%num_views]:.4f}', fontsize=10)
	else:
	axs[i].axis("off")
	plt.tight_layout()
	plt.savefig(save_path)
	plt.close(fig)
	print(f"Visualization file written to {save_path}")

	out_dir = save_path.replace('vis.png', 'reprojections')
	os.makedirs(out_dir, exist_ok=True)

	for i in range(num_views):
	gt = Image.fromarray(images[0, i])
	pred = Image.fromarray(images[1, i])
	gt.save(os.path.join(out_dir, f"gt_{i}.png"))
	pred.save(os.path.join(out_dir, f"pred_{i}.png"))

	return np.array(lpips_losses), np.array(mse_losses)