File size: 6,126 Bytes
d63ad23 d454202 d63ad23 1fa688f d63ad23 1fa688f 5777f44 1fa688f c6a8a22 9bb78c9 c6a8a22 2053232 c6a8a22 44c19b3 2053232 c6a8a22 83a8b53 f4e8cf6 fdeb0ad a14c9ce 1fa688f fdeb0ad 1fa688f a14c9ce 1fa688f fdeb0ad f4e8cf6 fdeb0ad 1fa688f f4e8cf6 1fa688f a14c9ce 1fa688f a14c9ce 1fa688f 76eeb7d 1fa688f f4e8cf6 a14c9ce 1fa688f 597a667 1fa688f c6a8a22 1fa688f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import numpy as np
import torch
import time
import nvdiffrast.torch as dr
from util.utils import get_tri
import tempfile
from mesh import Mesh
import zipfile
from util.renderer import Renderer
import trimesh
import xatlas
import cv2
from PIL import Image, ImageFilter
def vertex_color_to_uv_textured_glb(obj_path, glb_path, texture_size=512):
mesh = trimesh.load(obj_path, process=False)
vertex_colors = mesh.visual.vertex_colors[:, :3] # (N, 3), uint8
# Generate UVs
vmapping, indices, uvs = xatlas.parametrize(mesh.vertices, mesh.faces)
vertices = mesh.vertices[vmapping]
vertex_colors = vertex_colors[vmapping]
mesh.vertices = vertices
mesh.faces = indices
# Bake texture (hybrid: per-pixel barycentric for accuracy)
buffer_size = texture_size * 2
texture_buffer = np.zeros((buffer_size, buffer_size, 4), dtype=np.uint8)
face_uvs = uvs[mesh.faces]
face_colors = vertex_colors[mesh.faces]
min_xy = np.floor(np.min(face_uvs, axis=1) * (buffer_size - 1)).astype(int)
max_xy = np.ceil(np.max(face_uvs, axis=1) * (buffer_size - 1)).astype(int)
for i in range(len(mesh.faces)):
uv0, uv1, uv2 = face_uvs[i]
c0, c1, c2 = face_colors[i]
min_x, min_y = min_xy[i]
max_x, max_y = max_xy[i]
for y in range(min_y, max_y + 1):
for x in range(min_x, max_x + 1):
p = np.array([x + 0.5, y + 0.5]) / (buffer_size - 1)
# Barycentric coordinates
v0, v1, v2 = uv0, uv1, uv2
denom = (v1[1] - v2[1]) * (v0[0] - v2[0]) + (v2[0] - v1[0]) * (v0[1] - v2[1])
if denom == 0:
continue
u = ((v1[1] - v2[1]) * (p[0] - v2[0]) + (v2[0] - v1[0]) * (p[1] - v2[1])) / denom
v = ((v2[1] - v0[1]) * (p[0] - v2[0]) + (v0[0] - v2[0]) * (p[1] - v2[1])) / denom
w = 1 - u - v
if (u >= 0) and (v >= 0) and (w >= 0):
color = u * c0 + v * c1 + w * c2
texture_buffer[y, x, :3] = np.clip(color, 0, 255).astype(np.uint8)
texture_buffer[y, x, 3] = 255
# Inpainting, filtering, and downsampling (keep optimized)
image_bgra = texture_buffer.copy()
mask = (image_bgra[:, :, 3] == 0).astype(np.uint8) * 255
image_bgr = cv2.cvtColor(image_bgra, cv2.COLOR_BGRA2BGR)
inpainted_bgr = cv2.inpaint(image_bgr, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
inpainted_bgra = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2BGRA)
texture_buffer = inpainted_bgra[::-1]
image_texture = Image.fromarray(texture_buffer)
image_texture = image_texture.filter(ImageFilter.MedianFilter(size=3))
image_texture = image_texture.filter(ImageFilter.GaussianBlur(radius=1))
image_texture = image_texture.resize((texture_size, texture_size), Image.LANCZOS)
# Assign UVs and texture to mesh
material = trimesh.visual.material.PBRMaterial(
baseColorFactor=[1.0, 1.0, 1.0, 1.0],
baseColorTexture=image_texture,
metallicFactor=0.0,
roughnessFactor=1.0,
)
visuals = trimesh.visual.TextureVisuals(uv=uvs, material=material)
mesh.visual = visuals
mesh.export(glb_path)
image_texture.save("debug_texture.png")
def generate3d(model, rgb, ccm, device=None):
device = torch.device("cuda")
model.renderer = Renderer(tet_grid_size=model.tet_grid_size, camera_angle_num=model.camera_angle_num,
scale=model.input.scale, geo_type = model.geo_type)
color_tri = torch.from_numpy(rgb).to(device)/255
xyz_tri = torch.from_numpy(ccm[:,:,(2,1,0)]).to(device)/255
color = color_tri.permute(2,0,1)
xyz = xyz_tri.permute(2,0,1)
def get_imgs(color):
color_list = []
color_list.append(color[:,:,256*5:256*(1+5)])
for i in range(0,5):
color_list.append(color[:,:,256*i:256*(1+i)])
return torch.stack(color_list, dim=0)
triplane_color = get_imgs(color).permute(0,2,3,1).unsqueeze(0).to(device)
color = get_imgs(color)
xyz = get_imgs(xyz)
color = get_tri(color, dim=0, blender= True, scale = 1).unsqueeze(0).to(device)
xyz = get_tri(xyz, dim=0, blender= True, scale = 1, fix= True).unsqueeze(0).to(device)
triplane = torch.cat([color,xyz],dim=1).to(device)
model.eval()
if model.denoising == True:
tnew = 20
tnew = torch.randint(tnew, tnew+1, [triplane.shape[0]], dtype=torch.long, device=triplane.device)
noise_new = torch.randn_like(triplane) *0.5+0.5
triplane = model.scheduler.add_noise(triplane, noise_new, tnew)
start_time = time.time()
with torch.no_grad():
triplane_feature2 = model.unet2(triplane,tnew)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"unet takes {elapsed_time}s")
else:
triplane_feature2 = model.unet2(triplane)
with torch.no_grad():
data_config = {
'resolution': [1024, 1024],
"triview_color": triplane_color.to(device),
}
verts, faces = model.decode(data_config, triplane_feature2)
data_config['verts'] = verts[0]
data_config['faces'] = faces
from kiui.mesh_utils import clean_mesh
verts, faces = clean_mesh(data_config['verts'].squeeze().cpu().numpy().astype(np.float32), data_config['faces'].squeeze().cpu().numpy().astype(np.int32), repair = False, remesh=True, remesh_size=0.005, remesh_iters=1)
data_config['verts'] = torch.from_numpy(verts).to(device).contiguous()
data_config['faces'] = torch.from_numpy(faces).to(device).contiguous()
start_time = time.time()
with torch.no_grad():
mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f"", delete=False).name
model.export_mesh(data_config, mesh_path_glb, tri_fea_2 = triplane_feature2)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"uv takes {elapsed_time}s")
obj_path = mesh_path_glb + ".obj"
glb_path = mesh_path_glb + ".glb"
vertex_color_to_uv_textured_glb(obj_path, glb_path)
return glb_path |