implement postprocessing in torch

utils/postprocess.py

@@ -1,55 +1,45 @@
 import os
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = "3"
-import tensorflow as tf
 
 import sys
 import shutil
 import tqdm
 import random
 import math
+import torch
 import numpy as np
 import scipy as sp
 import scipy.io as sio
 import plotly.graph_objects as go
 import plotly.express as ex
+from pytorch3d.structures import Meshes
 
 from custom_utils.custom_utils import _maskFaces
-
-from scipy.spatial import cKDTree
-from tensorflow_graphics.geometry.representation.mesh.normals import vertex_normals
+from losses.losses import _nearest_neighbour
 
 
 # from hbertiche/PBNS
 """ COLLISION FUNCTIONS START """
-F = np.load("custom_utils/smpl_faces.npy")
+# F = np.load("smpl/smpl_faces.npy")
+to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
+F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
 
 
-def _nearest_neighbour(V, B):
-	tree = cKDTree(B.numpy())
-	return tree.query(V.numpy(), workers=-1)[1]
-
-@tf.function
-def _tf_nn_parallel(V, B):
-	return tf.map_fn(fn=lambda elem: tf.py_function(func=_nearest_neighbour, inp=elem, Tout=tf.int32), elems=(V,B), fn_output_signature=tf.int32, parallel_iterations=V.shape[0])
-
 
 def fix_collision(V, B, mask, thr=.004):
-    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
+    mesh = Meshes(B, F_fix_coll)
+    B_vn = torch.stack(mesh.verts_normals_list())
    
     # compute correspondences
-    idx = _tf_nn_parallel(V, B)
-    N, M = V.shape[:2]
-    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
-    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
-    idx = tf.stack((_idx, idx), axis=-1)
+    idx = _nearest_neighbour(V, B)
     # compute loss
-    D = V - tf.gather_nd(B, idx)
-    b_vn = tf.gather_nd(B_vn, idx)
-    dot = tf.einsum('abc,abc->ab', D, b_vn)
-    loss = tf.minimum(dot - thr, 0)  * mask
-    B_offset_closest = tf.gather_nd(B, idx) + (b_vn / 2e2)
-    output = tf.where(loss[:,:,tf.newaxis] < 0, B_offset_closest, V)
+    B_idx = torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
+    D = V - B_idx
+    b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
+    dot = torch.einsum('abc,abc->ab', D, b_vn)
+    loss = torch.minimum(dot - thr, torch.zeros_like(dot))  * mask
+    B_offset_closest = B_idx + (b_vn / 2e2)
+    output = torch.where(loss[:,:,None] < 0, B_offset_closest, V)
 
     return output