inference.py

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