infer_refine.py

import cv2
import numpy as np
import os
import trimesh
import argparse
import torch
import scipy
from PIL import Image

from refine.mesh_refine import geo_refine
from refine.func import make_star_cameras_orthographic
from refine.render import NormalsRenderer, calc_vertex_normals

from pytorch3d.structures import Meshes
from sklearn.neighbors import KDTree

from segment_anything import SamAutomaticMaskGenerator, sam_model_registry

sam = sam_model_registry["vit_h"](checkpoint="./ckpt/sam_vit_h_4b8939.pth").cuda()
generator = SamAutomaticMaskGenerator(
    model=sam,
    points_per_side=64,
    pred_iou_thresh=0.80,
    stability_score_thresh=0.92,
    crop_n_layers=1,
    crop_n_points_downscale_factor=2,
    min_mask_region_area=100,
)


def fix_vert_color_glb(mesh_path):
    from pygltflib import GLTF2, Material, PbrMetallicRoughness
    obj1 = GLTF2().load(mesh_path)
    obj1.meshes[0].primitives[0].material = 0
    obj1.materials.append(Material(
        pbrMetallicRoughness = PbrMetallicRoughness(
            baseColorFactor = [1.0, 1.0, 1.0, 1.0],
            metallicFactor = 0.,
            roughnessFactor = 1.0,
        ),
        emissiveFactor = [0.0, 0.0, 0.0],
        doubleSided = True,
    ))
    obj1.save(mesh_path)


def srgb_to_linear(c_srgb):
    c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4)
    return c_linear.clip(0, 1.)


def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
    # convert from pytorch3d meshes to trimesh mesh
    vertices = meshes.verts_packed().cpu().float().numpy()
    triangles = meshes.faces_packed().cpu().long().numpy()
    np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
    if save_glb_path.endswith(".glb"):
        # rotate 180 along +Y
        vertices[:, [0, 2]] = -vertices[:, [0, 2]]

    if apply_sRGB_to_LinearRGB:
        np_color = srgb_to_linear(np_color)
    assert vertices.shape[0] == np_color.shape[0]
    assert np_color.shape[1] == 3
    assert 0 <= np_color.min() and np_color.max() <= 1.001, f"min={np_color.min()}, max={np_color.max()}"
    np_color = np.clip(np_color, 0, 1)
    mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
    mesh.remove_unreferenced_vertices()
    # save mesh
    mesh.export(save_glb_path)
    if save_glb_path.endswith(".glb"):
        fix_vert_color_glb(save_glb_path)
    print(f"saving to {save_glb_path}")


def calc_horizontal_offset(target_img, source_img):
    target_mask = target_img.astype(np.float32).sum(axis=-1) > 750
    source_mask = source_img.astype(np.float32).sum(axis=-1) > 750
    best_offset = -114514
    for offset in range(-200, 200):
        offset_mask = np.roll(source_mask, offset, axis=1)
        overlap = (target_mask & offset_mask).sum()
        if overlap > best_offset:
            best_offset = overlap
            best_offset_value = offset
    return best_offset_value


def calc_horizontal_offset2(target_mask, source_img):
    source_mask = source_img.astype(np.float32).sum(axis=-1) > 750
    best_offset = -114514
    for offset in range(-200, 200):
        offset_mask = np.roll(source_mask, offset, axis=1)
        overlap = (target_mask & offset_mask).sum()
        if overlap > best_offset:
            best_offset = overlap
            best_offset_value = offset
    return best_offset_value


def get_distract_mask(color_0, color_1, normal_0=None, normal_1=None, thres=0.25, ratio=0.50, outside_thres=0.10, outside_ratio=0.20):
    distract_area = np.abs(color_0 - color_1).sum(axis=-1) > thres
    if normal_0 is not None and normal_1 is not None:
        distract_area |= np.abs(normal_0 - normal_1).sum(axis=-1) > thres
    labeled_array, num_features = scipy.ndimage.label(distract_area)
    results = []

    random_sampled_points = []

    for i in range(num_features + 1):
        if np.sum(labeled_array == i) > 1000 and np.sum(labeled_array == i) < 100000:
            results.append((i, np.sum(labeled_array == i)))
            # random sample a point in the area
            points = np.argwhere(labeled_array == i)
            random_sampled_points.append(points[np.random.randint(0, points.shape[0])])

    results = sorted(results, key=lambda x: x[1], reverse=True)  # [1:]
    distract_mask = np.zeros_like(distract_area)
    distract_bbox = np.zeros_like(distract_area)
    for i, _ in results:
        distract_mask |= labeled_array == i
        bbox = np.argwhere(labeled_array == i)
        min_x, min_y = bbox.min(axis=0)
        max_x, max_y = bbox.max(axis=0)
        distract_bbox[min_x:max_x, min_y:max_y] = 1

    points = np.array(random_sampled_points)[:, ::-1]
    labels = np.ones(len(points), dtype=np.int32)

    masks = generator.generate((color_1 * 255).astype(np.uint8))

    outside_area = np.abs(color_0 - color_1).sum(axis=-1) < outside_thres

    final_mask = np.zeros_like(distract_mask)
    for iii, mask in enumerate(masks):
        mask['segmentation'] = cv2.resize(mask['segmentation'].astype(np.float32), (1024, 1024)) > 0.5
        intersection = np.logical_and(mask['segmentation'], distract_mask).sum()
        total = mask['segmentation'].sum()
        iou = intersection / total
        outside_intersection = np.logical_and(mask['segmentation'], outside_area).sum()
        outside_total = mask['segmentation'].sum()
        outside_iou = outside_intersection / outside_total
        if iou > ratio and outside_iou < outside_ratio:
            final_mask |= mask['segmentation']

    # calculate coverage
    intersection = np.logical_and(final_mask, distract_mask).sum()
    total = distract_mask.sum()
    coverage = intersection / total

    if coverage < 0.8:
        # use original distract mask
        final_mask = (distract_mask.copy() * 255).astype(np.uint8)
        final_mask = cv2.dilate(final_mask, np.ones((3, 3), np.uint8), iterations=3)
        labeled_array_dilate, num_features_dilate = scipy.ndimage.label(final_mask)
        for i in range(num_features_dilate + 1):
            if np.sum(labeled_array_dilate == i) < 200:
                final_mask[labeled_array_dilate == i] = 255

        final_mask = cv2.erode(final_mask, np.ones((3, 3), np.uint8), iterations=3)
        final_mask = final_mask > 127

    return distract_mask, distract_bbox, random_sampled_points, final_mask


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_mv_dir', type=str, default='result/multiview')
    parser.add_argument('--input_obj_dir', type=str, default='result/slrm')
    parser.add_argument('--output_dir', type=str, default='result/refined')
    parser.add_argument('--outside_ratio', type=float, default=0.20)
    parser.add_argument('--no_decompose', action='store_true')
    args = parser.parse_args()

    for test_idx in os.listdir(args.input_mv_dir):
        mv_root_dir = os.path.join(args.input_mv_dir, test_idx)
        obj_dir = os.path.join(args.input_obj_dir, test_idx)

        fixed_v, fixed_f = None, None
        flow_vert, flow_vector = None, None
        last_colors, last_normals = None, None
        last_front_color, last_front_normal = None, None
        distract_mask = None

        mv, proj = make_star_cameras_orthographic(8, 1, r=1.2)
        mv = mv[[4, 3, 2, 0, 6, 5]]        
        renderer = NormalsRenderer(mv,proj,(1024,1024))

        if not args.no_decompose:  
            for name_idx, level in zip([3, 1, 2], [2, 1, 0]):
                mesh = trimesh.load(obj_dir + f'_{name_idx}.obj')
                new_mesh = mesh.split(only_watertight=False)
                new_mesh = [ j for j in new_mesh if len(j.vertices) >= 300 ]
                mesh = trimesh.Scene(new_mesh).dump(concatenate=True)
                mesh_v, mesh_f = mesh.vertices, mesh.faces

                if last_colors is None:
                    images = renderer.render(
                        torch.tensor(mesh_v, device='cuda').float(),
                        torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
                        torch.tensor(mesh_f, device='cuda'),
                    )
                    mask = (images[..., 3] < 0.9).cpu().numpy()

                colors, normals = [], []
                for i in range(6):
                    color_path = os.path.join(mv_root_dir, f'level{level}', f'color_{i}.png')
                    normal_path = os.path.join(mv_root_dir, f'level{level}', f'normal_{i}.png')
                    color = cv2.imread(color_path)
                    normal = cv2.imread(normal_path)
                    color = color[..., ::-1]
                    normal = normal[..., ::-1]

                    if last_colors is not None:
                        offset = calc_horizontal_offset(np.array(last_colors[i]), color)
                        # print('offset', i, offset)
                    else:
                        offset = calc_horizontal_offset2(mask[i], color)
                        # print('init offset', i, offset)

                    if offset != 0:
                        color = np.roll(color, offset, axis=1)
                        normal = np.roll(normal, offset, axis=1)

                    color = Image.fromarray(color)
                    normal = Image.fromarray(normal)
                    colors.append(color)
                    normals.append(normal)

                if last_front_color is not None and level == 0:
                    original_mask, distract_bbox, _, distract_mask = get_distract_mask(last_front_color, np.array(colors[0]).astype(np.float32) / 255.0, outside_ratio=args.outside_ratio)
                    cv2.imwrite(f'{args.output_dir}/{test_idx}/distract_mask.png', distract_mask.astype(np.uint8) * 255)
                else:  
                    distract_mask = None
                    distract_bbox = None

                last_front_color = np.array(colors[0]).astype(np.float32) / 255.0
                last_front_normal = np.array(normals[0]).astype(np.float32) / 255.0

                if last_colors is None:
                    from copy import deepcopy
                    last_colors, last_normals = deepcopy(colors), deepcopy(normals)

                # my mesh flow weight by nearest vertexs
                if fixed_v is not None and fixed_f is not None and level == 1:
                    t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)

                    fixed_v_cpu = fixed_v.cpu().numpy()
                    kdtree_anchor = KDTree(fixed_v_cpu)
                    kdtree_mesh_v = KDTree(mesh_v)
                    _, idx_anchor = kdtree_anchor.query(mesh_v, k=1)
                    _, idx_mesh_v = kdtree_mesh_v.query(mesh_v, k=25)
                    idx_anchor = idx_anchor.squeeze()
                    neighbors = torch.tensor(mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
                    # calculate the distances neighbors [V, 25, 3]; mesh_v [V, 3] -> [V, 25]
                    neighbor_dists = torch.norm(neighbors - torch.tensor(mesh_v).cuda()[:, None], dim=-1)
                    neighbor_dists[neighbor_dists > 0.06] = 114514.
                    neighbor_weights = torch.exp(-neighbor_dists * 1.)
                    neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
                    anchors = fixed_v[idx_anchor]  # V, 3
                    anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
                    dis_anchor = torch.clamp(((anchors - torch.tensor(mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
                    vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
                    vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
                    weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
                    mesh_v += weighted_vec_anchor.cpu().numpy()

                    t = trimesh.Trimesh(vertices=mesh_v, faces=mesh_f)

                mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
                mesh_f = torch.tensor(mesh_f, device='cuda')

                new_mesh, simp_v, simp_f = geo_refine(mesh_v, mesh_f, colors, normals, fixed_v=fixed_v, fixed_f=fixed_f, distract_mask=distract_mask, distract_bbox=distract_bbox)

                # my mesh flow weight by nearest vertexs
                try:
                    if fixed_v is not None and fixed_f is not None and level != 0:
                        new_mesh_v = new_mesh.verts_packed().cpu().numpy()

                        fixed_v_cpu = fixed_v.cpu().numpy()
                        kdtree_anchor = KDTree(fixed_v_cpu)
                        kdtree_mesh_v = KDTree(new_mesh_v)
                        _, idx_anchor = kdtree_anchor.query(new_mesh_v, k=1)
                        _, idx_mesh_v = kdtree_mesh_v.query(new_mesh_v, k=25)
                        idx_anchor = idx_anchor.squeeze()
                        neighbors = torch.tensor(new_mesh_v).cuda()[idx_mesh_v]  # V, 25, 3
                        # calculate the distances neighbors [V, 25, 3]; new_mesh_v [V, 3] -> [V, 25]
                        neighbor_dists = torch.norm(neighbors - torch.tensor(new_mesh_v).cuda()[:, None], dim=-1)
                        neighbor_dists[neighbor_dists > 0.06] = 114514.
                        neighbor_weights = torch.exp(-neighbor_dists * 1.)
                        neighbor_weights = neighbor_weights / neighbor_weights.sum(dim=1, keepdim=True)
                        anchors = fixed_v[idx_anchor]  # V, 3
                        anchor_normals = calc_vertex_normals(fixed_v, fixed_f)[idx_anchor]  # V, 3
                        dis_anchor = torch.clamp(((anchors - torch.tensor(new_mesh_v).cuda()) * anchor_normals).sum(-1), min=0) + 0.01
                        vec_anchor = dis_anchor[:, None] * anchor_normals  # V, 3
                        vec_anchor = vec_anchor[idx_mesh_v]  # V, 25, 3
                        weighted_vec_anchor = (vec_anchor * neighbor_weights[:, :, None]).sum(1)  # V, 3
                        new_mesh_v += weighted_vec_anchor.cpu().numpy()

                        # replace new_mesh verts with new_mesh_v
                        new_mesh = Meshes(verts=[torch.tensor(new_mesh_v, device='cuda')], faces=new_mesh.faces_list(), textures=new_mesh.textures)

                except Exception as e:
                    pass

                os.makedirs(f'{args.output_dir}/{test_idx}', exist_ok=True)
                save_py3dmesh_with_trimesh_fast(new_mesh, f'{args.output_dir}/{test_idx}/out_{level}.glb', apply_sRGB_to_LinearRGB=False)

                if fixed_v is None:
                    fixed_v, fixed_f = simp_v, simp_f
                else:
                    fixed_f = torch.cat([fixed_f, simp_f + fixed_v.shape[0]], dim=0)
                    fixed_v = torch.cat([fixed_v, simp_v], dim=0)


        else:
            mesh = trimesh.load(obj_dir + f'_0.obj')
            mesh_v, mesh_f = mesh.vertices, mesh.faces

            images = renderer.render(
                torch.tensor(mesh_v, device='cuda').float(),
                torch.ones_like(torch.from_numpy(mesh_v), device='cuda').float(),
                torch.tensor(mesh_f, device='cuda'),
            )
            mask = (images[..., 3] < 0.9).cpu().numpy()

            colors, normals = [], []
            for i in range(6):
                color_path = os.path.join(mv_root_dir, f'level0', f'color_{i}.png')
                normal_path = os.path.join(mv_root_dir, f'level0', f'normal_{i}.png')
                color = cv2.imread(color_path)
                normal = cv2.imread(normal_path)
                color = color[..., ::-1]
                normal = normal[..., ::-1]

                offset = calc_horizontal_offset2(mask[i], color)

                if offset != 0:
                    color = np.roll(color, offset, axis=1)
                    normal = np.roll(normal, offset, axis=1)

                color = Image.fromarray(color)
                normal = Image.fromarray(normal)
                colors.append(color)
                normals.append(normal)

            mesh_v = torch.tensor(mesh_v, device='cuda', dtype=torch.float32)
            mesh_f = torch.tensor(mesh_f, device='cuda')

            new_mesh, _, _ = geo_refine(mesh_v, mesh_f, colors, normals, no_decompose=True, expansion_weight=0.)

            os.makedirs(f'{args.output_dir}/{test_idx}', exist_ok=True)
            save_py3dmesh_with_trimesh_fast(new_mesh, f'{args.output_dir}/{test_idx}/out_nodecomp.glb', apply_sRGB_to_LinearRGB=False)