194 lines
7.6 KiB
Python
194 lines
7.6 KiB
Python
import torch.nn as nn
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import numpy as np
|
|
|
|
import cv2
|
|
import trimesh
|
|
import nvdiffrast.torch as dr
|
|
|
|
from CRM.model.archs.decoders.shape_texture_net import TetTexNet
|
|
from CRM.model.archs.unet import UNetPP
|
|
from CRM.util.renderer import Renderer
|
|
from CRM.model.archs.mlp_head import SdfMlp, RgbMlp
|
|
import xatlas
|
|
|
|
from mesh_processor.mesh import Mesh
|
|
|
|
|
|
class Dummy:
|
|
pass
|
|
|
|
class ConvolutionalReconstructionModel(nn.Module):
|
|
def __init__(self, specs):
|
|
super(ConvolutionalReconstructionModel, self).__init__()
|
|
|
|
self.specs = specs
|
|
# configs
|
|
input_specs = specs["Input"]
|
|
self.input = Dummy()
|
|
self.input.scale = input_specs['scale']
|
|
self.input.resolution = input_specs['resolution']
|
|
self.tet_grid_size = input_specs['tet_grid_size']
|
|
self.camera_angle_num = input_specs['camera_angle_num']
|
|
|
|
self.arch = Dummy()
|
|
self.arch.fea_concat = specs["ArchSpecs"]["fea_concat"]
|
|
self.arch.mlp_bias = specs["ArchSpecs"]["mlp_bias"]
|
|
|
|
self.dec = Dummy()
|
|
self.dec.c_dim = specs["DecoderSpecs"]["c_dim"]
|
|
self.dec.plane_resolution = specs["DecoderSpecs"]["plane_resolution"]
|
|
|
|
self.geo_type = specs["Train"].get("geo_type", "flex") # only support "flex"
|
|
|
|
self.unet2 = UNetPP(in_channels=self.dec.c_dim)
|
|
|
|
mlp_chnl_s = 3 if self.arch.fea_concat else 1 # 3 for queried triplane feature concatenation
|
|
self.decoder = TetTexNet(plane_reso=self.dec.plane_resolution, fea_concat=self.arch.fea_concat)
|
|
|
|
if self.geo_type == "flex":
|
|
self.weightMlp = nn.Sequential(
|
|
nn.Linear(mlp_chnl_s * 32 * 8, 512),
|
|
nn.SiLU(),
|
|
nn.Linear(512, 21))
|
|
|
|
self.sdfMlp = SdfMlp(mlp_chnl_s * 32, 512, bias=self.arch.mlp_bias)
|
|
self.rgbMlp = RgbMlp(mlp_chnl_s * 32, 512, bias=self.arch.mlp_bias)
|
|
self.renderer = Renderer(tet_grid_size=self.tet_grid_size, camera_angle_num=self.camera_angle_num,
|
|
scale=self.input.scale, geo_type = self.geo_type)
|
|
|
|
|
|
self.spob = True if specs['Pretrain']['mode'] is None else False # whether to add sphere
|
|
self.radius = specs['Pretrain']['radius'] # used when spob
|
|
|
|
self.denoising = True
|
|
from diffusers import DDIMScheduler
|
|
self.scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-2-1-base", subfolder="scheduler")
|
|
|
|
def decode(self, data, triplane_feature2):
|
|
if self.geo_type == "flex":
|
|
tet_verts = self.renderer.flexicubes.verts.unsqueeze(0)
|
|
tet_indices = self.renderer.flexicubes.indices
|
|
|
|
dec_verts = self.decoder(triplane_feature2, tet_verts)
|
|
out = self.sdfMlp(dec_verts)
|
|
|
|
weight = None
|
|
if self.geo_type == "flex":
|
|
grid_feat = torch.index_select(input=dec_verts, index=self.renderer.flexicubes.indices.reshape(-1),dim=1)
|
|
grid_feat = grid_feat.reshape(dec_verts.shape[0], self.renderer.flexicubes.indices.shape[0], self.renderer.flexicubes.indices.shape[1] * dec_verts.shape[-1])
|
|
weight = self.weightMlp(grid_feat)
|
|
weight = weight * 0.1
|
|
|
|
pred_sdf, deformation = out[..., 0], out[..., 1:]
|
|
if self.spob:
|
|
pred_sdf = pred_sdf + self.radius - torch.sqrt((tet_verts**2).sum(-1))
|
|
|
|
_, verts, faces = self.renderer(data, pred_sdf, deformation, tet_verts, tet_indices, weight= weight)
|
|
return verts[0].unsqueeze(0), faces[0].int()
|
|
|
|
def export_mesh(self, data, out_dir, ind, device=None, tri_fea_2 = None):
|
|
verts = data['verts']
|
|
faces = data['faces']
|
|
|
|
dec_verts = self.decoder(tri_fea_2, verts.unsqueeze(0))
|
|
colors = self.rgbMlp(dec_verts).squeeze().detach().cpu().numpy()
|
|
# Expect predicted colors value range from [-1, 1]
|
|
colors = (colors * 0.5 + 0.5).clip(0, 1)
|
|
|
|
verts = verts.squeeze().cpu().numpy()
|
|
faces = faces[..., [2, 1, 0]].squeeze().cpu().numpy()
|
|
|
|
# export the final mesh
|
|
with torch.no_grad():
|
|
mesh = trimesh.Trimesh(verts, faces, vertex_colors=colors, process=False) # important, process=True leads to seg fault...
|
|
mesh.export(out_dir / f'{ind}.obj')
|
|
|
|
@torch.no_grad()
|
|
def get_mesh_wt_uv(self, ctx, data, device, res, tri_fea_2=None):
|
|
|
|
mesh_v = data['verts'].squeeze().cpu().numpy()
|
|
mesh_pos_idx = data['faces'].squeeze().cpu().numpy()
|
|
|
|
def interpolate(attr, rast, attr_idx, rast_db=None):
|
|
return dr.interpolate(attr.contiguous(), rast, attr_idx, rast_db=rast_db,
|
|
diff_attrs=None if rast_db is None else 'all')
|
|
|
|
vmapping, indices, uvs = xatlas.parametrize(mesh_v, mesh_pos_idx)
|
|
|
|
mesh_v = torch.tensor(mesh_v, dtype=torch.float32, device=device)
|
|
mesh_pos_idx = torch.tensor(mesh_pos_idx, dtype=torch.int64, device=device)
|
|
|
|
# Convert to tensors
|
|
indices_int64 = indices.astype(np.uint64, casting='same_kind').view(np.int64)
|
|
|
|
uvs = torch.tensor(uvs, dtype=torch.float32, device=mesh_v.device)
|
|
mesh_tex_idx = torch.tensor(indices_int64, dtype=torch.int64, device=mesh_v.device)
|
|
# mesh_v_tex. ture
|
|
uv_clip = uvs[None, ...] * 2.0 - 1.0
|
|
|
|
# pad to four component coordinate
|
|
uv_clip4 = torch.cat((uv_clip, torch.zeros_like(uv_clip[..., 0:1]), torch.ones_like(uv_clip[..., 0:1])), dim=-1)
|
|
|
|
# rasterize
|
|
rast, _ = dr.rasterize(ctx, uv_clip4, mesh_tex_idx.int(), res)
|
|
|
|
# Interpolate world space position
|
|
gb_pos, _ = interpolate(mesh_v[None, ...], rast, mesh_pos_idx.int())
|
|
mask = rast[..., 3:4] > 0
|
|
|
|
# return uvs, mesh_tex_idx, gb_pos, mask
|
|
gb_pos_unsqz = gb_pos.view(-1, 3)
|
|
mask_unsqz = mask.view(-1)
|
|
tex_unsqz = torch.zeros_like(gb_pos_unsqz) + 1
|
|
|
|
gb_mask_pos = gb_pos_unsqz[mask_unsqz]
|
|
|
|
gb_mask_pos = gb_mask_pos[None, ]
|
|
|
|
with torch.no_grad():
|
|
|
|
dec_verts = self.decoder(tri_fea_2, gb_mask_pos)
|
|
colors = self.rgbMlp(dec_verts).squeeze()
|
|
|
|
# Expect predicted colors value range from [-1, 1]
|
|
lo, hi = (-1, 1)
|
|
colors = (colors - lo) * (255 / (hi - lo))
|
|
colors = colors.clip(0, 255)
|
|
|
|
tex_unsqz[mask_unsqz] = colors
|
|
|
|
tex = tex_unsqz.view(res + (3,))
|
|
|
|
verts = mesh_v.squeeze().cpu().numpy()
|
|
faces = mesh_pos_idx[..., [2, 1, 0]].squeeze().cpu().numpy()
|
|
# faces = mesh_pos_idx
|
|
# faces = faces.detach().cpu().numpy()
|
|
# faces = faces[..., [2, 1, 0]]
|
|
indices = indices[..., [2, 1, 0]]
|
|
|
|
verts = verts[:, [0, 2, 1]]
|
|
verts[:, 2] *= -1
|
|
|
|
#uvs[:, 1] = 1 - uvs[:, 1]
|
|
|
|
# Get albedo texture map
|
|
img = np.asarray(tex.data.cpu().numpy(), dtype=np.float32)
|
|
mask = np.sum(img.astype(float), axis=-1, keepdims=True)
|
|
mask = (mask <= 3.0).astype(float)
|
|
kernel = np.ones((3, 3), 'uint8')
|
|
dilate_img = cv2.dilate(img, kernel, iterations=1)
|
|
img = img * (1 - mask) + dilate_img * mask
|
|
img = img.clip(0, 255).astype(np.uint8) / 255
|
|
|
|
mesh_v = torch.tensor(verts, dtype=torch.float32, device=device)
|
|
mesh_f = torch.tensor(faces, dtype=torch.int32, device=device)
|
|
mesh_vt = uvs.clone().detach()
|
|
mesh_ft = torch.tensor(indices.astype(np.int32), dtype=torch.int32, device=device)
|
|
albedo = torch.tensor(img, dtype=torch.float32, device=device)
|
|
|
|
mesh = Mesh(v=mesh_v, f=mesh_f, vt=mesh_vt, ft=mesh_ft, albedo=albedo, device=device)
|
|
mesh.auto_normal()
|
|
return mesh |