Refactor device handling across multiple modules to enforce CUDA usage explicitly. This change ensures consistent device allocation for model operations, enhancing performance on systems with GPU support while maintaining compatibility.
fdeb0ad
| 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=640): | |
| 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 | |
| buffer_size = texture_size * 2 | |
| texture_buffer = np.zeros((buffer_size, buffer_size, 4), dtype=np.uint8) | |
| def barycentric_interpolate(v0, v1, v2, c0, c1, c2, p): | |
| v0v1 = v1 - v0 | |
| v0v2 = v2 - v0 | |
| v0p = p - v0 | |
| d00 = np.dot(v0v1, v0v1) | |
| d01 = np.dot(v0v1, v0v2) | |
| d11 = np.dot(v0v2, v0v2) | |
| d20 = np.dot(v0p, v0v1) | |
| d21 = np.dot(v0p, v0v2) | |
| denom = d00 * d11 - d01 * d01 | |
| if abs(denom) < 1e-8: | |
| return (c0 + c1 + c2) / 3 | |
| v = (d11 * d20 - d01 * d21) / denom | |
| w = (d00 * d21 - d01 * d20) / denom | |
| u = 1.0 - v - w | |
| u = np.clip(u, 0, 1) | |
| v = np.clip(v, 0, 1) | |
| w = np.clip(w, 0, 1) | |
| return u * c0 + v * c1 + w * c2 | |
| def is_point_in_triangle(p, v0, v1, v2): | |
| def sign(p1, p2, p3): | |
| return (p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1]) | |
| d1 = sign(p, v0, v1) | |
| d2 = sign(p, v1, v2) | |
| d3 = sign(p, v2, v0) | |
| has_neg = (d1 < 0) or (d2 < 0) or (d3 < 0) | |
| has_pos = (d1 > 0) or (d2 > 0) or (d3 > 0) | |
| return not (has_neg and has_pos) | |
| for face in mesh.faces: | |
| uv0, uv1, uv2 = uvs[face] | |
| c0, c1, c2 = vertex_colors[face] | |
| uv0 = (uv0 * (buffer_size - 1)).astype(int) | |
| uv1 = (uv1 * (buffer_size - 1)).astype(int) | |
| uv2 = (uv2 * (buffer_size - 1)).astype(int) | |
| min_x = max(int(np.floor(min(uv0[0], uv1[0], uv2[0]))), 0) | |
| max_x = min(int(np.ceil(max(uv0[0], uv1[0], uv2[0]))), buffer_size - 1) | |
| min_y = max(int(np.floor(min(uv0[1], uv1[1], uv2[1]))), 0) | |
| max_y = min(int(np.ceil(max(uv0[1], uv1[1], uv2[1]))), buffer_size - 1) | |
| 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]) | |
| if is_point_in_triangle(p, uv0, uv1, uv2): | |
| color = barycentric_interpolate(uv0, uv1, uv2, c0, c1, c2, p) | |
| texture_buffer[y, x, :3] = np.clip(color, 0, 255).astype(np.uint8) | |
| texture_buffer[y, x, 3] = 255 | |
| # Inpainting, filtering, and downsampling | |
| 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 |