monoarti/mesh_utils.py

import numpy as np
import torch
import torch.nn.functional as F
import collections
from collections import defaultdict
import cv2
import random
import math
import quaternion
from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh import TexturesVertex, TexturesUV, Textures
import os
import shutil
import imageio
from typing import Optional, List


def triangulate_pcd(sem_mask, semantic_pred, render_size, border_size):
    #verts = torch.ones(render_size).nonzero().numpy()
    verts = sem_mask.nonzero().numpy()
    vert_id_map = defaultdict(dict)
    for idx, vert in enumerate(verts):
        vert_id_map[vert[0]][vert[1]] = idx# + len(verts)

    height = render_size[0] - border_size * 2
    width = render_size[1] - border_size * 2

    semantic_pred = semantic_pred.numpy()

    triangles = []
    for vert in verts:
        # upper right triangle
        if (
            vert[0] < height - 1
            and vert[1] < width - 1
            and sem_mask[vert[0] + 1][vert[1] + 1]
            and sem_mask[vert[0]][vert[1] + 1]
            and semantic_pred[vert[0]][vert[1]] == semantic_pred[vert[0] + 1][vert[1] + 1]
            and semantic_pred[vert[0]][vert[1]] == semantic_pred[vert[0]][vert[1] + 1]
        ):
            triangles.append(
                [
                    vert_id_map[vert[0]][vert[1]],
                    vert_id_map[vert[0] + 1][vert[1] + 1],
                    vert_id_map[vert[0]][vert[1] + 1],
                ]
            )
        # bottom left triangle
        if (
            vert[0] < height - 1
            and vert[1] < width - 1
            and sem_mask[vert[0] + 1][vert[1] + 1]
            and sem_mask[vert[0]][vert[1] + 1]
            and semantic_pred[vert[0]][vert[1]] == semantic_pred[vert[0] + 1][vert[1]]
            and semantic_pred[vert[0]][vert[1]] == semantic_pred[vert[0] + 1][vert[1] + 1]
        ):
            triangles.append(
                [
                    vert_id_map[vert[0]][vert[1]],
                    vert_id_map[vert[0] + 1][vert[1]],
                    vert_id_map[vert[0] + 1][vert[1] + 1],
                ]
            )
    triangles = np.array(triangles)
    triangles = torch.LongTensor(triangles)

    return triangles


def get_pcd(verts, normal, offset, h=480, w=640, focal_length=517.97):
    """
    Copy from 
    https://github.com/JasonQSY/Articulation3D/blob/master/articulation3d/articulation3d/utils/vis.py

    convert 2d verts to 3d point cloud based on plane normal and offset
    depth = offset / n \dot K^{-1}q
    """
    offset_x = w/2
    offset_y = h/2
    K = [[focal_length, 0, offset_x],
        [0, focal_length, offset_y],
        [0, 0, 1]]
    K_inv = np.linalg.inv(np.array(K))
    homogeneous = np.hstack((verts, np.ones(len(verts)).reshape(-1,1)))
    ray = K_inv@homogeneous.T
    depth = offset / np.dot(normal, ray)
    pcd = depth.reshape(-1,1) * ray.T

    #import pdb; pdb.set_trace()
    return pcd


def fit_homography(src_pts, tgt_pts):
    """
    Fit a homography from src_pts to tgt_pts.

    src_pts: torch.LongTensor shape (N x 2)
    tgt_pts: torch.LongTensor shape (N x 2)
    
    """
    src_pts = src_pts.numpy().astype(np.float32)
    tgt_pts = tgt_pts.numpy().astype(np.float32)

    N = 4
    
    # randomly pick up 4 control points
    #ids = random.sample(range(src_pts.shape[0]), 4)
    #src_pts = 
    #import pdb; pdb.set_trace()
    H, mask = cv2.findHomography(src_pts, tgt_pts, cv2.RANSAC, 5.0)
    #H = cv2.getPerspectiveTransform(obj_mask.nonzero().cpu().numpy().astype(np.float32), pts_reproj[1].numpy().astype(np.float32))

    return H

def create_cylinder_mesh(radius, p0, p1, stacks=10, slices=10):

    def compute_length_vec3(vec3):
        return math.sqrt(vec3[0]*vec3[0] + vec3[1]*vec3[1] + vec3[2]*vec3[2])
    
    def rotation(axis, angle):
        rot = np.eye(4)
        c = np.cos(-angle)
        s = np.sin(-angle)
        t = 1.0 - c
        axis /= compute_length_vec3(axis)
        x = axis[0]
        y = axis[1]
        z = axis[2]
        rot[0,0] = 1 + t*(x*x-1)
        rot[0,1] = z*s+t*x*y
        rot[0,2] = -y*s+t*x*z
        rot[1,0] = -z*s+t*x*y
        rot[1,1] = 1+t*(y*y-1)
        rot[1,2] = x*s+t*y*z
        rot[2,0] = y*s+t*x*z
        rot[2,1] = -x*s+t*y*z
        rot[2,2] = 1+t*(z*z-1)
        return rot


    verts = []
    indices = []
    diff = (p1 - p0).astype(np.float32)
    height = compute_length_vec3(diff)
    for i in range(stacks+1):
        for i2 in range(slices):
            theta = i2 * 2.0 * math.pi / slices
            pos = np.array([radius*math.cos(theta), radius*math.sin(theta), height*i/stacks])
            verts.append(pos)
    for i in range(stacks):
        for i2 in range(slices):
            i2p1 = math.fmod(i2 + 1, slices)
            indices.append( np.array([(i + 1)*slices + i2, i*slices + i2, i*slices + i2p1], dtype=np.uint32) )
            indices.append( np.array([(i + 1)*slices + i2, i*slices + i2p1, (i + 1)*slices + i2p1], dtype=np.uint32) )
    transform = np.eye(4)
    va = np.array([0, 0, 1], dtype=np.float32)
    vb = diff
    vb /= compute_length_vec3(vb)
    axis = np.cross(vb, va)
    angle = np.arccos(np.clip(np.dot(va, vb), -1, 1))
    if angle != 0:
        if compute_length_vec3(axis) == 0:
            dotx = va[0]
            if (math.fabs(dotx) != 1.0):
                axis = np.array([1,0,0]) - dotx * va
            else:
                axis = np.array([0,1,0]) - va[1] * va
            axis /= compute_length_vec3(axis)
        transform = rotation(axis, -angle)
    transform[:3,3] += p0
    verts = [np.dot(transform, np.array([v[0], v[1], v[2], 1.0])) for v in verts]
    verts = [np.array([v[0], v[1], v[2]]) / v[3] for v in verts]
        
    return verts, indices


def create_arrow_mesh(radius, p0, p1, stacks=10, slices=10, arrow_height=0):

    def compute_length_vec3(vec3):
        return math.sqrt(vec3[0]*vec3[0] + vec3[1]*vec3[1] + vec3[2]*vec3[2])
    
    def rotation(axis, angle):
        rot = np.eye(4)
        c = np.cos(-angle)
        s = np.sin(-angle)
        t = 1.0 - c
        axis /= compute_length_vec3(axis)
        x = axis[0]
        y = axis[1]
        z = axis[2]
        rot[0,0] = 1 + t*(x*x-1)
        rot[0,1] = z*s+t*x*y
        rot[0,2] = -y*s+t*x*z
        rot[1,0] = -z*s+t*x*y
        rot[1,1] = 1+t*(y*y-1)
        rot[1,2] = x*s+t*y*z
        rot[2,0] = y*s+t*x*z
        rot[2,1] = -x*s+t*y*z
        rot[2,2] = 1+t*(z*z-1)
        return rot


    verts = []
    indices = []
    diff = (p1 - p0).astype(np.float32)
    height = compute_length_vec3(diff)
    for i in range(stacks+2):
        if i == stacks+1:
            # arrow tip
            cur_radius = 0
            cur_height = height
        elif i == stacks:
            # arrow base
            cur_radius = radius*3
            cur_height = height * (1-arrow_height) * (i-1)/stacks
        else:
            # cylinder
            cur_radius = radius
            cur_height = height * (1-arrow_height) * i/stacks
        for i2 in range(slices):
            theta = i2 * 2.0 * math.pi / slices
            pos = np.array([cur_radius*math.cos(theta), cur_radius*math.sin(theta), cur_height])
            verts.append(pos)
    for i in range(stacks+1):
        for i2 in range(slices):
            i2p1 = math.fmod(i2 + 1, slices)
            indices.append( np.array([(i + 1)*slices + i2, i*slices + i2, i*slices + i2p1], dtype=np.uint32) )
            indices.append( np.array([(i + 1)*slices + i2, i*slices + i2p1, (i + 1)*slices + i2p1], dtype=np.uint32) )
    transform = np.eye(4)
    va = np.array([0, 0, 1], dtype=np.float32)
    vb = diff
    vb /= compute_length_vec3(vb)
    axis = np.cross(vb, va)
    angle = np.arccos(np.clip(np.dot(va, vb), -1, 1))
    if angle != 0:
        if compute_length_vec3(axis) == 0:
            dotx = va[0]
            if (math.fabs(dotx) != 1.0):
                axis = np.array([1,0,0]) - dotx * va
            else:
                axis = np.array([0,1,0]) - va[1] * va
            axis /= compute_length_vec3(axis)
        transform = rotation(axis, -angle)
    transform[:3,3] += p0
    verts = [np.dot(transform, np.array([v[0], v[1], v[2], 1.0])) for v in verts]
    verts = [np.array([v[0], v[1], v[2]]) / v[3] for v in verts]
        
    return verts, indices



def get_camera_meshes(camera_list, radius=0.02):
    verts_list = []
    faces_list = []
    color_list = []
    rots = np.array([quaternion.as_rotation_matrix(camera_info['rotation']) for camera_info in camera_list])

    # ai habitat frame
    lookat = np.array([0,0,-1])
    vertical = np.array([0,1,0])

    positions = np.array([camera_info['position'] for camera_info in camera_list])
    lookats = rots@lookat.T
    verticals = rots@vertical.T
    predetermined_color = [
        [0.10196, 0.32157, 1.0],
        [1.0, 0.0667, 0.1490],# [0.8314, 0.0667, 0.3490],
        # [0.0, 0.4392156862745098, 0.7529411764705882],
        # [0.3764705882352941, 0.08627450980392155, 0.47843137254901963],
    ]
    for idx, (position, lookat, vertical, color) in enumerate(zip(positions, lookats, verticals, predetermined_color)):
        cur_num_verts = 0
        # r, g, b = create_color_palette()[idx+10]
        edges = get_cone_edges(position, lookat, vertical)
        # color = [r/255.0,g/255.0,b/255.0]
        cam_verts = []
        cam_inds = []
        for k in range(len(edges)):
            cyl_verts, cyl_ind = create_cylinder_mesh(radius, edges[k][0], edges[k][1])
            cyl_verts = [x for x in cyl_verts]
            cyl_ind = [x + cur_num_verts for x in cyl_ind]
            cur_num_verts += len(cyl_verts)
            cam_verts.extend(cyl_verts)
            cam_inds.extend(cyl_ind)
        # Create a textures object
        verts_list.append(torch.tensor(cam_verts, dtype=torch.float32))
        faces_list.append(torch.tensor(cam_inds, dtype=torch.float32))
        color_list.append(color)

    color_tensor = torch.tensor(color_list, dtype=torch.float32).unsqueeze_(1)
    #tex = Textures(verts_uvs=None, faces_uvs=None, verts_rgb=color_tensor)
    tex = TexturesVertex(verts_features=color_tensor)

    # Initialise the mesh with textures
    meshes = Meshes(verts=verts_list, faces=faces_list, textures=tex)
    return meshes


def get_cone_edges(position, lookat, vertical):
    def get_cone_verts(position, lookat, vertical):
        vertical = np.array(vertical) / np.linalg.norm(vertical)
        lookat = np.array(lookat) / np.linalg.norm(lookat)
        right = np.cross(np.array(lookat), np.array(vertical))
        right = right / np.linalg.norm(right)
        top = np.cross(right, lookat)
        top = top / np.linalg.norm(top)

        right *= .4
        lookat *= .4
        top *= .1
        verts = {
            'topR': position + lookat + top + right,
            'topL': position + lookat + top - right,
            'center': position,
            'bottomR': position + lookat - top + right,
            'bottomL': position + lookat - top - right,
        }
        return verts

    cone_verts = get_cone_verts(position, lookat, vertical)
    edges = [
        (cone_verts['center'], cone_verts['topR']),
        (cone_verts['center'], cone_verts['topL']),
        (cone_verts['center'], cone_verts['bottomR']),
        (cone_verts['center'], cone_verts['bottomL']),
        (cone_verts['topR'], cone_verts['topL']),
        (cone_verts['bottomR'], cone_verts['topR']),
        (cone_verts['bottomR'], cone_verts['bottomL']),
        (cone_verts['topL'], cone_verts['bottomL']),
    ]
    return edges




def get_axis_mesh(radius, pt1, pt2):
    verts_list = []
    faces_list = []
    color_list = []
    cyl_verts, cyl_ind = create_arrow_mesh(radius, pt1.numpy(), pt2.numpy())
    cyl_verts = [x for x in cyl_verts]
    cyl_ind = [x for x in cyl_ind]

    # Create a textures object
    verts_list.append(torch.tensor(cyl_verts, dtype=torch.float32))
    faces_list.append(torch.tensor(cyl_ind, dtype=torch.float32))
    # color_list.append([0.10196, 0.32157, 1.0])
    # color_tensor = torch.tensor(color_list, dtype=torch.float32).unsqueeze_(1)
    # Textures(verts_uvs=axis_verts_rgb, faces_uvs=axis_pt1.faces_list(), maps=torch.zeros((1,5,5,3)).cuda())
    # tex = TexturesVertex(verts_features=color_tensor)

    # Initialise the mesh with textures
    meshes = Meshes(verts=verts_list, faces=faces_list)
    return meshes


def save_obj_articulation(folder, prefix, meshes, cam_meshes=None, decimal_places=None, blend_flag=False, map_files=None, uv_maps=None):
    os.makedirs(folder, exist_ok=True)

    # pytorch3d does not support map_files
    #map_files = meshes.textures.map_files()
    #assert map_files is not None
    if map_files is None and uv_maps is None:
        raise RuntimeError("either map_files or uv_maps should be set!")

    # generate map_files from uv_map
    if uv_maps is not None and map_files is None:
        map_files = []
        uv_dir = os.path.join(folder, 'uv_maps')
        if not os.path.exists(uv_dir):
            os.mkdir(uv_dir)
        for map_id, uv_map in enumerate(uv_maps):
            uv_path = os.path.join(uv_dir, '{}_uv_plane_{}.png'.format(prefix, map_id))
            #pdb.set_trace()
            imageio.imwrite(uv_path, uv_map)
            map_files.append(uv_path)

    #pdb.set_trace()

    f_mtl = open(os.path.join(folder, prefix+'.mtl'), 'w')
    f = open(os.path.join(folder, prefix+'.obj'), 'w')
    try:
        seen = set()
        uniq_map_files = [m for m in list(map_files) if m not in seen and not seen.add(m)]
        for map_id, map_file in enumerate(uniq_map_files):
            if uv_maps is not None:
                # we do not need to copy map_files,
                # they are already in uv_maps/...
                f_mtl.write(_get_mtl_map(
                    os.path.basename(map_file).split('.')[0], 
                    os.path.join('uv_maps', os.path.basename(map_file))
                ))
                continue

            if not blend_flag:
                shutil.copy(map_file, folder)
                os.chmod(os.path.join(folder, os.path.basename(map_file)), 0o755)
                f_mtl.write(_get_mtl_map(os.path.basename(map_file).split('.')[0], os.path.basename(map_file)))
            else:
                rgb = cv2.imread(map_file, cv2.IMREAD_COLOR)
                if cam_meshes is not None:
                    blend_color = np.array(cam_meshes.textures.verts_features_packed().numpy().tolist()[map_id])*255
                else:
                    blend_color = np.array(create_color_palette()[map_id+10])
                alpha = 0.7
                blend = (rgb*alpha + blend_color[::-1]*(1-alpha)).astype(np.uint8)
                cv2.imwrite(os.path.join(folder, os.path.basename(map_file).split('.')[0]+'_debug.png'), blend)
                f_mtl.write(_get_mtl_map(os.path.basename(map_file).split('.')[0], os.path.basename(map_file).split('.')[0]+'_debug.png'))
        
        f.write(f"mtllib {prefix}.mtl\n\n")
        # we want [list]    verts, vert_uvs, map_files; 
        #         [packed]  faces;
        #         face per mesh
        verts_list = meshes.verts_list()
        verts_uvs_list = meshes.textures.verts_uvs_list()
        faces_list = meshes.faces_packed().split(meshes.num_faces_per_mesh().tolist(), dim=0)
        #pdb.set_trace()
        for idx, (verts, verts_uvs, faces, map_file) in enumerate(zip(verts_list, verts_uvs_list, faces_list, map_files)):
            f.write(f"# mesh {idx}\n")
            trunc_verts_uvs = verts_uvs[:verts.shape[0]]
            _save(f, verts, faces, verts_uv=trunc_verts_uvs, map_file=map_file, idx=idx, decimal_places=decimal_places)
        if cam_meshes:
            face_offset = np.sum([len(v) for v in verts_list])
            cam_verts_list = cam_meshes.verts_list()
            cam_verts_rgbs_list = cam_meshes.textures.verts_features_packed().numpy().tolist()
            cam_faces_list = (cam_meshes.faces_packed()+face_offset).split(cam_meshes.num_faces_per_mesh().tolist(), dim=0)
            assert(len(cam_verts_rgbs_list) == len(cam_verts_list))
            for idx, (verts, faces, rgb) in enumerate(zip(cam_verts_list, cam_faces_list, cam_verts_rgbs_list)):
                f.write(f"# camera {idx}\n")
                f_mtl.write(_get_mtl_rgb(idx, rgb))
                _save(f, verts, faces, rgb=rgb, idx=idx, decimal_places=decimal_places)
    finally:
        f.close()
        f_mtl.close()


def _get_mtl_map(material_name, map_Kd):
        return f"""newmtl {material_name}
map_Kd {map_Kd}
# Test colors
Ka 1.000 1.000 1.000  # white
Kd 1.000 1.000 1.000  # white
Ks 0.000 0.000 0.000  # black
Ns 10.0\n"""


def _get_mtl_rgb(material_idx, rgb):
        return f"""newmtl color_{material_idx}
Kd {rgb[0]} {rgb[1]} {rgb[2]}
Ka 0.000 0.000 0.000\n"""


def _save(f, verts, faces, verts_uv=None, map_file=None, rgb=None, idx=None, double_sided=True, decimal_places: Optional[int] = None):
    if decimal_places is None:
        float_str = "%f"
    else:
        float_str = "%" + ".%df" % decimal_places

    lines = ""
    
    V, D = verts.shape
    for i in range(V):
        vert = [float_str % verts[i, j] for j in range(D)]
        lines += "v %s\n" % " ".join(vert)

    if verts_uv is not None:
        V, D = verts_uv.shape
        for i in range(V):
            vert_uv = [float_str % verts_uv[i, j] for j in range(D)]
            lines += "vt %s\n" % " ".join(vert_uv)

    if map_file is not None:
        lines += f"usemtl {os.path.basename(map_file).split('.')[0]}\n"
    elif rgb is not None:
        lines += f"usemtl color_{idx}\n"    

    if faces != []:
        F, P = faces.shape
        for i in range(F):
            if verts_uv is not None:
                face = ["%d/%d" % (faces[i, j] + 1, faces[i, j] + 1) for j in range(P)]
            else:
                face = ["%d" % (faces[i, j] + 1) for j in range(P)]
            # if i + 1 < F:
            lines += "f %s\n" % " ".join(face)
            if double_sided:
                if verts_uv is not None:
                    face = ["%d/%d" % (faces[i, j] + 1, faces[i, j] + 1) for j in reversed(range(P))]
                else:
                    face = ["%d" % (faces[i, j] + 1) for j in reversed(range(P))]
                lines += "f %s\n" % " ".join(face)
            # elif i + 1 == F:
            #     # No newline at the end of the file.
            #     lines += "f %s" % " ".join(face)
    else:
        print(f"face = []")
    f.write(lines)