utils.py

import numpy as np
import torch
import cv2
import os

def saveObj(filename, materialName, vertices, faces, normals = None, tcoords = None, textureFileName = 'texture.png'):
    '''
    write mesh to an obj file
    :param filename: path to where to save the obj file
    :param materialFileName: material name
    :param vertices:  float tensor [n, 3]
    :param faces: tensor [#triangles, 3]
    :param normals: float tensor [n, 3]
    :param tcoords: float tensor [n, 2]
    :param textureFileName: name of the texture to use with material
    :return:
    '''
    assert(vertices.dim() == 2 and  vertices.shape[-1] == 3)
    assert (faces.dim() == 2 and faces.shape[-1] == 3)

    if normals is not None:
        assert (normals.dim() == 2 and normals.shape[-1] == 3)

    if tcoords is not None:
        assert (tcoords.dim() == 2 and tcoords.shape[-1] == 2)

    if torch.is_tensor(vertices):
        vertices = vertices.detach().cpu().numpy()
    if torch.is_tensor(faces):
        faces = faces.detach().cpu().numpy()
    if torch.is_tensor(normals):
        normals = normals.detach().cpu().numpy()
    if torch.is_tensor(tcoords):
        tcoords = tcoords.detach().cpu().numpy()

    assert(isinstance(vertices, np.ndarray))
    assert (isinstance(faces, np.ndarray))
    assert (isinstance(normals, np.ndarray))
    assert (isinstance(tcoords, np.ndarray))

    #write material
    f = open(os.path.dirname(filename) + '/' + materialName, 'w')
    f.write('newmtl material0\n')
    f.write('map_Kd ' + textureFileName + '\n')
    f.close()

    f = open(filename, 'w')
    f.write('###########################################################\n')
    f.write('# OBJ file generated by faceYard 2021\n')
    f.write('#\n')
    f.write('# Num Vertices: %d\n' % (vertices.shape[0]))
    f.write('# Num Triangles: %d\n' % (faces.shape[0]))
    f.write('#\n')
    f.write('###########################################################\n')
    f.write('\n')
    f.write('mtllib ' + materialName + '\n')

    #write vertices
    for v in vertices:
        f.write('v %f %f %f\n' % (v[0], v[1], v[2]))

    # write the tcoords
    if tcoords is not None and tcoords.shape[0] > 0:
        for uv in tcoords:
            f.write('vt %f %f\n' % (uv[0], uv[1]))

    #write the normals
    if normals is not None and normals.shape[0] > 0:
        for n in normals:
            f.write('vn %f %f %f\n' % (n[0], n[1], n[2]))

    f.write('usemtl material0\n')
    #write face indices list
    for t in faces:
        f.write('f %d/%d/%d %d/%d/%d %d/%d/%d\n' % (t[0] + 1, t[0] + 1,t[0] + 1,
                                              t[1] + 1, t[1] + 1,t[1] + 1,
                                              t[2] + 1, t[2] + 1, t[2] + 1))
    f.close()
def saveLandmarksVerticesProjections(imageTensor, projPoints, landmarks):
    '''
    for debug, render the projected vertices and landmakrs on image
    :param images: [w, h, 3]
    :param projPoints: [n, 3]
    :param landmarks: [n, 2]
    :return: tensor [w, h, 3
    '''
    assert(imageTensor.dim() == 3 and imageTensor.shape[-1] == 3 )
    assert(projPoints.dim() == 2 and projPoints.shape[-1] == 2)
    assert(projPoints.shape == landmarks.shape)
    image = imageTensor.clone().detach().cpu().numpy() * 255.
    landmarkCount = landmarks.shape[0]
    for i in range(landmarkCount):
        x = landmarks[i, 0]
        y = landmarks[i, 1]
        cv2.circle(image, (int(x), int(y)), 2, (0, 255, 0), -1)
        x = projPoints[i, 0]
        y = projPoints[i, 1]
        cv2.circle(image, (int(x), int(y)), 2, (0, 0, 255), -1)

    return image

def mkdir_p(path):
    import errno
    import os

    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise
def loadDictionaryFromPickle(picklePath):
    import pickle
    handle = open(picklePath, 'rb')
    assert handle is not None
    dic = pickle.load(handle)
    handle.close()
    return dic
def writeDictionaryToPickle(dict, picklePath):
    import pickle
    handle = open(picklePath, 'wb')
    pickle.dump(dict, handle, pickle.HIGHEST_PROTOCOL)
    handle.close()