3D_Stage/refine.py

import os
import numpy as np
import torch
import nvdiffrast.torch as dr
import json
import torch.nn.functional as F
from PIL import Image
import pymeshlab
import cv2

def back_to_texture(glctx, look_at, pos, tri, tex, uv, uv_idx, idx, vn):
    rast_out, rast_out_db = dr.rasterize(glctx, pos, tri, resolution=[tex.shape[0],tex.shape[1]])
    gb_normal, _ = dr.interpolate(vn[None], rast_out, tri)
    gb_normal = F.normalize(gb_normal, dim=-1)
    if idx == 2 or idx == 0:
        filter_camera = [torch.tensor([[[[1,0.,0.]]]]).cuda(), torch.tensor([[[[-1,0.,0.]]]]).cuda()]
    else:
        filter_camera = [torch.tensor([[[[0,-1.,0.]]]]).cuda(), torch.tensor([[[[0,1.,0.]]]]).cuda()]
    nmasks = []
    for fc in filter_camera:
        nmasks.append(((gb_normal * fc) > 0.75).int().sum(keepdim=True, dim=-1))
    gb_normal_mask = 1 - (nmasks[0] | nmasks[1])
   #Image.fromarray(np.clip(gb_normal_mask[0,...,0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"mask_normal_{idx}.png")
    gb_mask = rast_out[...,3:4] > 0
    tri_list = torch.unique(rast_out[...,3:4].reshape(-1))
    tri_list = (tri_list[1:] - 1).to(torch.int32)
    pos = pos[0]

    depth_map = rast_out[...,3:4].clone()
    depth_map[depth_map > 0] = 1
    depth_map = depth_map.to(torch.float32)
    dmax = (rast_out[...,2:3] * gb_mask).max()
    uv = torch.cat([uv * 2 - 1, torch.zeros(uv.shape[0], 1).cuda(), torch.ones(uv.shape[0], 1).cuda()], dim=1).unsqueeze(0)
    uv_idx = uv_idx[tri_list.to(torch.long)]
    rast_uv, rast_uv_db = dr.rasterize(glctx, uv, uv_idx, resolution=(1024, 1024))
    pos_clip = torch.cat([pos[...,:2], pos[...,3:]], -1)
    pos_2d, _ = dr.interpolate(pos_clip, rast_uv, tri[tri_list.to(torch.long)]) # pos (x, y, z, w)
    pos_coord = (pos_2d[...,:2] / (pos_2d[...,2:3] + 1e-6) + 1) / 2.
    texture_mask = (rast_uv[...,3:4] > 0).int()
    color = dr.texture(tex[None, ...] * gb_normal_mask, pos_coord, filter_mode='linear')
    color_mask = dr.texture(gb_normal_mask.to(torch.float32), pos_coord, filter_mode='linear')
    color_mask[color_mask > 0.82] = 1
    color_mask[color_mask <= 0.82] = 0
    color_mask = color_mask.to(torch.int32)
   #Image.fromarray(np.clip(color_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"depth_{idx}.png")
    texture_mask = texture_mask * color_mask
   #Image.fromarray(np.clip(color[0].detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).save(f"{idx}.png")
   #Image.fromarray(np.clip(texture_mask[0].repeat(1,1,3).detach().cpu().numpy() * 255, 0, 255).astype(np.uint8)).convert("RGB").save(f"mask-{idx}.png")
    return color, texture_mask, rast_uv

def perspective(fovy=0.6913, aspect=1.0, n=0.1, f=1000.0, device=None):
    y = np.tan(fovy / 2)
    return torch.tensor([[1/(y*aspect),    0,            0,              0], 
                         [           0, 1/-y,            0,              0], 
                         [           0,    0, -(f+n)/(f-n), -(2*f*n)/(f-n)], 
                         [           0,    0,           -1,              0]]).to(torch.float32).cuda()

def rec_mvp(trans, h, w):
    mv = trans
    fov = 40. / 180. * np.pi
    proj = perspective(fov, h / w, n=0.1, f=1000)
    mvp = proj @ mv
    return mvp

def aggregate_texture(kd_map, textures, texture_masks, rast_uvs):
    texture = torch.zeros_like(textures[0])
    texture_mask = torch.zeros_like(texture_masks[0])
    ctex = []
    for idx in range(len(textures)):
        ctex.append(textures[idx] * texture_masks[idx] + 10 * (1 - texture_masks[idx]))
    cat_textures = torch.stack(ctex, dim=-2)
    dis_measure = (cat_textures - kd_map.unsqueeze(-2)).abs().sum(-1)
    _, choose_idx = dis_measure.min(-1)

    choose_idx = choose_idx.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 1, 1, 3)
    final_texture_map = torch.gather(cat_textures, 3, choose_idx).squeeze(-2)
    #cv2.imwrite("final_texture_map.png", cv2.cvtColor((final_texture_map[0].cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    #cv2.imwrite("final_texture_mask.png", (texture_mask[0].cpu().numpy() * 255).astype(np.uint8))
    zero_mask = (final_texture_map.max(dim=-1, keepdim=True)[0] > 0.1)
    close_mask = ((final_texture_map[0] - kd_map).abs().sum(dim=-1, keepdim=True) < 1.0).int()
    for idx in range(len(textures)):
        texture += textures[idx] * texture_masks[idx]
        texture_mask |= texture_masks[idx]
    texture_mask = texture_mask * zero_mask * close_mask[None]
    optimize_mask = (texture_mask == 0).int()

   #import pdb; pdb.set_trace()
   #mask = (texture_mask[0].cpu().numpy() * 255).astype(np.uint8)
   #cv2.imwrite("mask.png", mask)
   #kernel = np.ones((5,5), np.uint8)
   #dilated = cv2.dilate(mask, kernel, iterations=1)
   #cv2.imwrite("di_mask.png", dilated)
   #texture_mask[0] = torch.from_numpy(dilated).unsqueeze(-1).to(torch.float32) / 255.

    final_texture_map = final_texture_map[0] * texture_mask[0]
    Image.fromarray(np.rint(final_texture_map.cpu().numpy() * 255).astype(np.uint8)).save(f"final_texture.png")

   #cv2.imwrite("kd_map.png", cv2.cvtColor((kd_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
   #cv2.imwrite("texture_map.png", cv2.cvtColor((final_texture_map.cpu().numpy() * 255).astype(np.uint8), cv2.COLOR_BGR2RGB))
   #result = cv2.seamlessClone((final_texture_map.cpu().numpy() * 255).astype(np.uint8), (kd_map.cpu().numpy() * 255).astype(np.uint8), mask, (mask.shape[1]//2, mask.shape[0]//2), cv2.NORMAL_CLONE)
   #cv2.imwrite("result.png", cv2.cvtColor(result * 255, cv2.COLOR_BGR2RGB))

    kd_map = kd_map * (1 - texture_mask[0]) + final_texture_map
    return kd_map, optimize_mask

def refine(save_path, front_image, back_image, left_image, right_image):
    ms = pymeshlab.MeshSet()
    mesh_path = f"{save_path}/model-00.obj"
    ms.load_new_mesh(mesh_path)
    ms.apply_coord_laplacian_smoothing(stepsmoothnum=10)
    tl = open(mesh_path, "r").readlines()
    tex_uv = []
    uv_idx = []
    for line in tl:
        if line.startswith("vt"):
            uvs = line.split(" ")[1:3]
            tex_uv += [float(uvs[0]), 1.0-float(uvs[1])]
    tex_uv = torch.from_numpy(np.array(tex_uv)).to(torch.float32).cuda().reshape(-1, 2)
    m = ms.current_mesh()
    v_matrix = m.vertex_matrix()
    f_matrix = m.face_matrix()
    vn = m.vertex_normal_matrix()
    uv_idx = torch.arange(f_matrix.shape[0] * 3).reshape(-1, 3).to(torch.int32).cuda()
    vn = torch.tensor(vn).contiguous().cuda().to(torch.float32)

    frames = []
    front_camera = torch.tensor([[
        1,0,0,0,
        0,0,1,0,
        0,-1,0,-1.5,
        0,0,0,1,
    ]]).to(torch.float32).reshape(4,4).cuda()
    back_camera = torch.tensor([[
        1,0,0,0,
        0,0,1,0,
        0,1,0,-1.5,
        0,0,0,1,
    ]]).to(torch.float32).reshape(4,4).cuda()
    right_camera = torch.tensor([[
        0,-1,0,0,
        0,0,1,0,
        1,0,0,-1.5,
        0,0,0,1,
    ]]).to(torch.float32).reshape(4,4).cuda()
    left_camera = torch.tensor([[
        0,1,0,0,
        0,0,1,0,
        -1,0,0,-1.5,
        0,0,0,1,
    ]]).to(torch.float32).reshape(4,4).cuda()
    frames = [front_camera, left_camera, back_camera, right_camera]

    target_images = []
    for target_image in [front_image, left_image, back_image, right_image]:
        target_images.append(torch.from_numpy(np.asarray(target_image.convert("RGB"))).to(torch.float32).cuda() / 255.)

    pos = torch.tensor(v_matrix, dtype=torch.float32).contiguous().cuda()
    tri = torch.tensor(f_matrix, dtype=torch.int32).contiguous().cuda()

    kd_map = (torch.tensor(np.asarray(Image.open(f"{save_path}/texture_kd.jpg"))) / 255.).cuda()
    translate_tensor = torch.zeros((1,1,3)).cuda()
    pos = torch.cat([pos, torch.ones([pos.shape[0], 1]).cuda()],-1).unsqueeze(0)
    glctx = dr.RasterizeCudaContext()
    target_texture = []
    target_mask = []
    rast_uvs = []
    with torch.no_grad():
        for idx, trans in enumerate(frames):
            target_image = target_images[idx]
            look_at = -torch.linalg.inv(trans)[:3,2]
            mvp = rec_mvp(trans, h=target_images[0].shape[0], w=target_images[0].shape[1])
            trans_pos = pos.clone()
            trans_pos[...,:3] += translate_tensor
            view_pos = torch.matmul(mvp, trans_pos.unsqueeze(-1)).squeeze(-1) 
            texture, mask, rast_uv = back_to_texture(glctx, look_at, view_pos, tri, target_image, tex_uv, uv_idx, idx, vn)
            target_texture.append(texture)
            target_mask.append(mask)
            rast_uvs.append(rast_uv)
        kd_map, opt_mask = aggregate_texture(kd_map, target_texture, target_mask, rast_uvs)
        opt_mask = opt_mask[0]
    Image.fromarray((np.clip(kd_map.detach().cpu().numpy() * 255, 0, 255)).astype(np.uint8)).save(f"{save_path}/refined_texture_kd.jpg")

   #ms.save_current_mesh(f"{save_path}/model-00.obj")
    with open(f"{save_path}/model-00.mtl", "w") as f:
        f.write(f"newmtl default\nKa 0.0 0.0 0.0\nmap_Kd refined_texture_kd.jpg\nKs 0.0 0.0 0.0")