refine/mesh_refine.py

import torch
import spaces
import numpy as np
import trimesh
from PIL import Image
from typing import List
from tqdm import tqdm
from sklearn.neighbors import KDTree

from refine.func import from_py3d_mesh, get_cameras_list, make_star_cameras_orthographic, multiview_color_projection, simple_clean_mesh, to_py3d_mesh, to_pyml_mesh
from refine.opt import MeshOptimizer
from refine.render import NormalsRenderer, calc_vertex_normals

import pytorch3d
from pytorch3d.structures import Meshes
import xatlas
import cv2


def mesh_uv_wrap(vertices, faces):
    if len(faces) > 50000:
        raise ValueError("The mesh has more than 50,000 faces, which is not supported.")

    vmapping, indices, uvs = xatlas.parametrize(vertices, faces)
    return vertices[vmapping], indices, uvs


def stride_from_shape(shape):
    stride = [1]
    for x in reversed(shape[1:]):
        stride.append(stride[-1] * x)
    return list(reversed(stride))

def scatter_add_nd_with_count(input, count, indices, values, weights=None):
    # input: [..., C], D dimension + C channel
    # count: [..., 1], D dimension
    # indices: [N, D], long
    # values: [N, C]

    D = indices.shape[-1]
    C = input.shape[-1]
    size = input.shape[:-1]
    stride = stride_from_shape(size)

    assert len(size) == D

    input = input.view(-1, C)  # [HW, C]
    count = count.view(-1, 1)

    flatten_indices = (indices * torch.tensor(stride,
                                              dtype=torch.long, device=indices.device)).sum(-1)  # [N]

    if weights is None:
        weights = torch.ones_like(values[..., :1])

    input.scatter_add_(0, flatten_indices.unsqueeze(1).repeat(1, C), values)
    count.scatter_add_(0, flatten_indices.unsqueeze(1), weights)

    return input.view(*size, C), count.view(*size, 1)


def linear_grid_put_2d(H, W, coords, values, return_count=False):
    # coords: [N, 2], float in [0, 1]
    # values: [N, C]

    C = values.shape[-1]

    indices = coords * torch.tensor(
        [H - 1, W - 1], dtype=torch.float32, device=coords.device
    )
    indices_00 = indices.floor().long()  # [N, 2]
    indices_00[:, 0].clamp_(0, H - 2)
    indices_00[:, 1].clamp_(0, W - 2)
    indices_01 = indices_00 + torch.tensor(
        [0, 1], dtype=torch.long, device=indices.device
    )
    indices_10 = indices_00 + torch.tensor(
        [1, 0], dtype=torch.long, device=indices.device
    )
    indices_11 = indices_00 + torch.tensor(
        [1, 1], dtype=torch.long, device=indices.device
    )

    h = indices[..., 0] - indices_00[..., 0].float()
    w = indices[..., 1] - indices_00[..., 1].float()
    w_00 = (1 - h) * (1 - w)
    w_01 = (1 - h) * w
    w_10 = h * (1 - w)
    w_11 = h * w

    result = torch.zeros(H, W, C, device=values.device,
                         dtype=values.dtype)  # [H, W, C]
    count = torch.zeros(H, W, 1, device=values.device,
                        dtype=values.dtype)  # [H, W, 1]
    weights = torch.ones_like(values[..., :1])  # [N, 1]

    result, count = scatter_add_nd_with_count(
        result, count, indices_00, values * w_00.unsqueeze(1), weights * w_00.unsqueeze(1))
    result, count = scatter_add_nd_with_count(
        result, count, indices_01, values * w_01.unsqueeze(1), weights * w_01.unsqueeze(1))
    result, count = scatter_add_nd_with_count(
        result, count, indices_10, values * w_10.unsqueeze(1), weights * w_10.unsqueeze(1))
    result, count = scatter_add_nd_with_count(
        result, count, indices_11, values * w_11.unsqueeze(1), weights * w_11.unsqueeze(1))

    if return_count:
        return result, count

    mask = (count.squeeze(-1) > 0)
    result[mask] = result[mask] / count[mask].repeat(1, C)

    return result, count.squeeze(-1) == 0


def remove_color(arr):
    if arr.shape[-1] == 4:
        arr = arr[..., :3]
    # calc diffs
    base = arr[0, 0]
    diffs = np.abs(arr.astype(np.int32) - base.astype(np.int32)).sum(axis=-1)
    alpha = (diffs <= 80)
    
    arr[alpha] = 255
    alpha = ~alpha
    arr = np.concatenate([arr, alpha[..., None].astype(np.int32) * 255], axis=-1)
    return arr


def simple_remove(imgs):
    """Only works for normal"""
    if not isinstance(imgs, list):
        imgs = [imgs]
        single_input = True
    else:
        single_input = False
    rets = []
    for img in imgs:
        arr = np.array(img)
        arr = remove_color(arr)
        rets.append(Image.fromarray(arr.astype(np.uint8)))
    if single_input:
        return rets[0]
    return rets


def erode_alpha(img_list):
    out_img_list = []
    for idx, img in enumerate(img_list):
        arr = np.array(img)
        alpha = (arr[:, :, 3] > 127).astype(np.uint8)
        # erode 1px
        import cv2
        alpha = cv2.erode(alpha, np.ones((3, 3), np.uint8), iterations=1)
        alpha = (alpha * 255).astype(np.uint8)
        img = Image.fromarray(np.concatenate([arr[:, :, :3], alpha[:, :, None]], axis=-1))
        out_img_list.append(img)
    return out_img_list


def merge_small_faces(mesh, thres=1e-5):
    area_faces = mesh.area_faces
    small_faces = area_faces < thres

    vertices = mesh.vertices
    faces = mesh.faces

    new_vertices = vertices.tolist()
    vertex_mapping = {}
    
    for face_idx in np.where(small_faces)[0]:
        face = faces[face_idx]
        v1, v2, v3 = face
        center = np.mean(vertices[face], axis=0)

        new_vertex_idx = len(new_vertices)
        new_vertices.append(center)

        vertex_mapping[v1] = new_vertex_idx
        vertex_mapping[v2] = new_vertex_idx
        vertex_mapping[v3] = new_vertex_idx

    for k,v in vertex_mapping.items():
        faces[faces == k] = v

    faces = faces[~small_faces]

    new_mesh = trimesh.Trimesh(vertices=new_vertices, faces=faces, postprocess=False)
    new_mesh.remove_unreferenced_vertices()
    new_mesh.remove_degenerate_faces()
    new_mesh.remove_duplicate_faces()
    
    return new_mesh


def init_target(img_pils, new_bkgd=(0., 0., 0.), device="cuda"):
    # Convert the background color to a PyTorch tensor
    new_bkgd = torch.tensor(new_bkgd, dtype=torch.float32).view(1, 1, 3).to(device)
    
    # Convert all images to PyTorch tensors and process them
    imgs = torch.stack([torch.from_numpy(np.array(img, dtype=np.float32)) for img in img_pils]).to(device) / 255
    img_nps = imgs[..., :3]
    alpha_nps = imgs[..., 3]
    ori_bkgds = img_nps[:, :1, :1]
    
    # Avoid divide by zero and calculate the original image
    alpha_nps_clamp = torch.clamp(alpha_nps, 1e-6, 1)
    ori_img_nps = (img_nps - ori_bkgds * (1 - alpha_nps.unsqueeze(-1))) / alpha_nps_clamp.unsqueeze(-1)
    ori_img_nps = torch.clamp(ori_img_nps, 0, 1)
    img_nps = torch.where(alpha_nps.unsqueeze(-1) > 0.05, ori_img_nps * alpha_nps.unsqueeze(-1) + new_bkgd * (1 - alpha_nps.unsqueeze(-1)), new_bkgd)

    rgba_img_np = torch.cat([img_nps, alpha_nps.unsqueeze(-1)], dim=-1)
    return rgba_img_np


def reconstruct_stage1(pils: List[Image.Image], steps=100, vertices=None, faces=None, fixed_v=None, fixed_f=None, lr=0.03, start_edge_len=0.15, end_edge_len=0.005,
                       decay=0.995, loss_expansion_weight=0.1, gain=0.1, remesh_interval=1, remesh_start=0, distract_mask=None, distract_bbox=None):
    vertices, faces = vertices.cuda(), faces.cuda()
    assert len(pils) == 6
    mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
    mv = mv[[4, 3, 2, 0, 6, 5]]
    renderer = NormalsRenderer(mv,proj,list(pils[0].size))

    target_images = init_target(pils, new_bkgd=(0., 0., 0.))

    # init from coarse mesh
    opt = MeshOptimizer(vertices, faces, local_edgelen=False, gain=gain, edge_len_lims=(end_edge_len, start_edge_len), lr=lr,
                        remesh_interval=remesh_interval, remesh_start=remesh_start)

    _vertices = opt.vertices
    _faces = opt.faces

    if fixed_v is not None and fixed_f is not None:
        kdtree = KDTree(fixed_v.cpu().numpy())

    mask = target_images[..., -1] < 0.5

    for i in tqdm(range(steps)):
        faces = torch.cat([_faces, fixed_f + len(_vertices)], dim=0) if fixed_f is not None else _faces
        vertices = torch.cat([_vertices, fixed_v], dim=0) if fixed_v is not None else _vertices

        opt.zero_grad()
        opt._lr *= decay
        normals = calc_vertex_normals(vertices,faces)

        normals[:, 0] *= -1
        normals[:, 2] *= -1

        images = renderer.render(vertices,normals,faces)
        loss_expand = 0.5 * ((vertices+normals).detach() - vertices).pow(2).mean()
        
        t_mask = images[..., -1] > 0.5
        loss_target_l2 = (images[t_mask] - target_images[t_mask]).abs().pow(2).mean()
        loss_alpha_target_mask_l2 = (images[..., -1][mask] - target_images[..., -1][mask]).pow(2).mean()
        
        loss = loss_target_l2 + loss_alpha_target_mask_l2 + loss_expand * loss_expansion_weight

        if distract_mask is not None:
            hair_visible_normals = normals
            hair_visible_normals[len(_vertices):] = -1.
            _images = renderer.render(vertices,hair_visible_normals,faces)
            loss_distract = (_images[0][distract_mask] - target_images[0][distract_mask]).pow(2).mean()

            target_outside = target_images[0][..., :3].clone()
            target_outside[~distract_mask] = 0.

            loss_outside_distract = (_images[0][..., :3][~distract_mask] - target_outside[..., :3][~distract_mask]).pow(2).mean()

            loss = loss + loss_distract * 1. + loss_outside_distract * 10.

        if fixed_v is not None and fixed_f is not None:
            _, idx = kdtree.query(_vertices.detach().cpu().numpy(), k=1)
            idx = idx.squeeze()
            anchors = fixed_v[idx].detach()

            normals_fixed = calc_vertex_normals(fixed_v, fixed_f)
            loss_anchor = (torch.clamp(((anchors - _vertices) * normals_fixed[idx]).sum(-1), min=-0)+0).pow(3)
            loss_anchor_dist_mask = (anchors - _vertices).norm(dim=-1) < 0.05
            loss_anchor = loss_anchor[loss_anchor_dist_mask].mean()

            loss = loss + loss_anchor * 100.
        
        # out of box
        loss_oob = (vertices.abs() > 0.99).float().mean() * 10
        loss = loss + loss_oob

        loss.backward()
        opt.step()

        if i % remesh_interval == 0 and i >= remesh_start:
            _vertices,_faces = opt.remesh(poisson=False)

    vertices, faces = opt._vertices.detach(), opt._faces.detach()
    
    return vertices, faces


def run_mesh_refine(vertices, faces, pils: List[Image.Image], fixed_v=None, fixed_f=None, steps=100, start_edge_len=0.02, end_edge_len=0.005,
                    decay=0.99, update_normal_interval=10, update_warmup=10, return_mesh=True, process_inputs=True, process_outputs=True, remesh_interval=20):
    poission_steps = []

    assert len(pils) == 6
    mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
    mv = mv[[4, 3, 2, 0, 6, 5]]        
    renderer = NormalsRenderer(mv,proj,list(pils[0].size))

    target_images = init_target(pils, new_bkgd=(0., 0., 0.)) # 4s

    # init from coarse mesh
    opt = MeshOptimizer(vertices, faces, ramp=5, edge_len_lims=(end_edge_len, start_edge_len), local_edgelen=False, laplacian_weight=0.02)

    _vertices = opt.vertices
    _faces = opt.faces
    alpha_init = None

    mask = target_images[..., -1] < 0.5

    for i in tqdm(range(steps)):
        faces = torch.cat([_faces, fixed_f + len(_vertices)], dim=0) if fixed_f is not None else _faces
        vertices = torch.cat([_vertices, fixed_v], dim=0) if fixed_v is not None else _vertices

        opt.zero_grad()
        opt._lr *= decay
        normals = calc_vertex_normals(vertices,faces)
        images = renderer.render(vertices,normals,faces)
        if alpha_init is None:
            alpha_init = images.detach()
        
        if i < update_warmup or i % update_normal_interval == 0:
            with torch.no_grad():
                py3d_mesh = to_py3d_mesh(vertices, faces, normals)
                cameras = get_cameras_list(azim_list = [180, 225, 270, 0, 90, 135], device=vertices.device, focal=1/1.2)
                _, _, target_normal = from_py3d_mesh(multiview_color_projection(py3d_mesh, pils, cameras_list=cameras, weights=[2,0.8,0.8,2,0.8,0.8], confidence_threshold=0.1, complete_unseen=False, below_confidence_strategy='original', reweight_with_cosangle='linear'))
                target_normal = target_normal * 2 - 1
                target_normal = torch.nn.functional.normalize(target_normal, dim=-1)

                target_normal[:, 0] *= -1
                target_normal[:, 2] *= -1

                debug_images = renderer.render(vertices,target_normal,faces)
        
        d_mask = images[..., -1] > 0.5
        loss_debug_l2 = (images[..., :3][d_mask] - debug_images[..., :3][d_mask]).pow(2).mean()
        
        loss_alpha_target_mask_l2 = (images[..., -1][mask] - target_images[..., -1][mask]).pow(2).mean()
        
        loss = loss_debug_l2 + loss_alpha_target_mask_l2
        
        # out of box
        loss_oob = (vertices.abs() > 0.99).float().mean() * 10
        loss = loss + loss_oob
        
        loss.backward()
        opt.step()
        
        if i % remesh_interval == 0:
            _vertices,_faces = opt.remesh(poisson=(i in poission_steps))

    vertices, faces = opt._vertices.detach(), opt._faces.detach()
    
    if process_outputs:
        vertices = vertices / 2 * 1.35
        vertices[..., [0, 2]] = - vertices[..., [0, 2]]

    return vertices, faces


def geo_refine(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=None, fixed_f=None,
               distract_mask=None, distract_bbox=None, thres=3e-6, no_decompose=False):
    print(mesh_v.device, mesh_f.device)
    if fixed_v is not None:
        print('fixed_v', fixed_v.shape, fixed_v.device)
    if fixed_f is not None:
        print('fixed_f', fixed_f.shape, fixed_f.device)
    vertices, faces = geo_refine_1(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=expansion_weight, fixed_v=fixed_v, fixed_f=fixed_f,
                 distract_mask=distract_mask, distract_bbox=distract_bbox, thres=thres, no_decompose=no_decompose)
    vertices, faces = geo_refine_2(vertices, faces, fixed_v=fixed_v)
    return geo_refine_3(vertices, faces, rgb_ls, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask)

def geo_refine_1(mesh_v, mesh_f, rgb_ls, normal_ls, expansion_weight=0.1, fixed_v=None, fixed_f=None,
                 distract_mask=None, distract_bbox=None, thres=3e-6, no_decompose=False):
    
    rm_normals = simple_remove(normal_ls)

    # transfer the alpha channel of rm_normals to img_list
    for idx, img in enumerate(rm_normals):
        rgb_ls[idx] = Image.fromarray(np.concatenate([np.array(rgb_ls[idx])[..., :3], np.array(img)[:, :, 3:4]], axis=-1))
    assert np.mean(np.array(rgb_ls[0])[..., 3]) < 250
    
    rgb_ls = erode_alpha(rgb_ls)

    stage1_lr = 0.08 if fixed_v is None else 0.01
    stage1_remesh_interval = 1 if fixed_v is None else 30

    if no_decompose:
        stage1_lr = 0.03
        stage1_remesh_interval = 30
    
    if fixed_v is not None:
        return mesh_v, mesh_f

    vertices, faces = reconstruct_stage1(rm_normals, steps=200, vertices=mesh_v, faces=mesh_f, fixed_v=fixed_v, fixed_f=fixed_f,
                                         lr=stage1_lr, remesh_interval=stage1_remesh_interval, start_edge_len=0.04,
                                         end_edge_len=0.010, gain=0.05, loss_expansion_weight=expansion_weight,
                                         distract_mask=distract_mask, distract_bbox=distract_bbox)

    vertices, faces = run_mesh_refine(vertices, faces, rm_normals, fixed_v=fixed_v, fixed_f=fixed_f, steps=100, start_edge_len=0.010, end_edge_len=0.001,
                                      decay=0.99, update_normal_interval=20, update_warmup=5, process_inputs=False, process_outputs=False, remesh_interval=1)
    return vertices, faces

def geo_refine_2(vertices, faces, fixed_v=None):
    meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=2, apply_sub_divide=False, sub_divide_threshold=0.25)
    simp_vertices, simp_faces = meshes.verts_packed(), meshes.faces_packed()
    vertices, faces = simp_vertices.detach().cpu().numpy(), simp_faces.detach().cpu().numpy()
    # vertices, faces = trimesh.remesh.subdivide(vertices, faces)
    return vertices, faces

def geo_refine_3(vertices_, faces_, rgb_ls, fixed_v=None, fixed_f=None, distract_mask=None):
    # vertices, faces, uvs = mesh_uv_wrap(vertices_, faces_)
    vmapping, indices, uvs = xatlas.parametrize(vertices_, faces_)
    vertices, faces = vertices_[vmapping], indices
    
    def subdivide(vertices, faces, uvs):
        vertices, faces = trimesh.remesh.subdivide(
            vertices=np.hstack((vertices, uvs.copy())),
            faces=faces
        )
        return vertices[:, :3], faces, vertices[:, 3:]

    if fixed_v is not None:
        dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs = subdivide(vertices, faces, uvs)
        dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs = subdivide(dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs)
        # dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs = subdivide(dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs)
        dense_vertices, dense_faces = trimesh.remesh.subdivide(vertices_, faces_)
        dense_vertices, dense_faces = trimesh.remesh.subdivide(dense_vertices, dense_faces)
        # dense_vertices, dense_faces = trimesh.remesh.subdivide(dense_vertices, dense_faces)
    else:
        dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs = subdivide(vertices, faces, uvs)
        dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs = subdivide(dense_atlas_vertices, dense_atlas_faces, dense_atlas_uvs)
        dense_vertices, dense_faces = trimesh.remesh.subdivide(vertices_, faces_)
        dense_vertices, dense_faces = trimesh.remesh.subdivide(dense_vertices, dense_faces)

    origin_len_v, origin_len_f = len(dense_vertices), len(dense_faces)

    # concatenate fixed_v and fixed_f
    if fixed_v is not None and fixed_f is not None:
        dense_vertices, dense_faces = np.concatenate([dense_vertices, fixed_v.detach().cpu().numpy()], axis=0), np.concatenate([dense_faces, fixed_f.detach().cpu().numpy() + len(dense_vertices)], axis=0)
    dense_vertices, dense_faces = torch.from_numpy(dense_vertices).cuda(), torch.from_numpy(dense_faces.astype('int32')).cuda()
    # reconstruct meshes
    meshes = Meshes(verts=[dense_vertices], faces=[dense_faces], textures=pytorch3d.renderer.mesh.textures.TexturesVertex([torch.zeros_like(dense_vertices).float()]))
    new_meshes = multiview_color_projection(meshes, rgb_ls, resolution=1024, device="cuda", complete_unseen=True, confidence_threshold=0.2, cameras_list = get_cameras_list([180, 225, 270, 0, 90, 135], "cuda", focal=1/1.2), weights=[2.0, 0.5, 0.0, 1.0, 0.0, 0.5] if distract_mask is None else [2.0, 0.0, 0.5, 1.0, 0.5, 0.0], distract_mask=distract_mask)
    
    if fixed_v is not None and fixed_f is not None:
        dense_vertices = dense_vertices[:origin_len_v]
        dense_faces = dense_faces[:origin_len_f]
        textures = new_meshes.textures.verts_features_packed()[:origin_len_v]
    else:
        textures = new_meshes.textures.verts_features_packed()

    # distances = torch.cdist(torch.tensor(dense_atlas_vertices).cuda(), torch.tensor(dense_vertices).cuda())
    # nearest_indices = torch.argmin(distances, dim=1)
    # atlas_textures = textures[nearest_indices]

    chunk_size = 500
    atlas_textures_chunks = []
    for i in range(0, len(dense_atlas_vertices), chunk_size):
        chunk = dense_atlas_vertices[i:i+chunk_size]
        distances = torch.cdist(torch.tensor(chunk).cuda(), torch.tensor(dense_vertices).cuda())
        nearest_indices = torch.argmin(distances, dim=1)
        atlas_textures_chunks.append(textures[nearest_indices])
    atlas_textures = torch.cat(atlas_textures_chunks, dim=0)

    dense_atlas_uvs = torch.tensor(dense_atlas_uvs, dtype=torch.float32).cuda()
    tex_img, mask = linear_grid_put_2d(1024, 1024, dense_atlas_uvs, atlas_textures)
    tex_img, mask = tex_img.cpu().numpy(), mask.cpu().numpy()
    tex_img = cv2.inpaint((tex_img * 255).astype(np.uint8), (mask*255).astype('uint8'), 3, cv2.INPAINT_NS)
    tex_img = Image.fromarray(np.transpose(tex_img,(1,0,2))[::-1])

    mesh = trimesh.Trimesh(vertices, faces, process=False)
    # material = trimesh.visual.texture.SimpleMaterial(image=tex_img, diffuse=(255, 255, 255))
    material = trimesh.visual.material.PBRMaterial(
        roughnessFactor=1.0,
        baseColorTexture=tex_img,
        baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8)
    )
    texture_visuals = trimesh.visual.TextureVisuals(uv=uvs, image=tex_img, material=material)
    mesh.visual = texture_visuals

    return mesh, torch.tensor(vertices).cuda(), torch.tensor(faces.astype('int64')).cuda()