Try to fix the gpu aborted problem : processing on downsampled meshes during optimization

app.py

@@ -28,6 +28,7 @@ import yaml
 from utils.surfaces import Surface
 import notebook_helpers as helper
 from utils.meshplot import visu_pts
+from utils.fmap import FM_to_p2p
 from utils.torch_fmap import extract_p2p_torch_fmap, torch_zoomout
 import torch
 import argparse
@@ -94,7 +95,7 @@ DEFAULT_SETTINGS = {
     "deepfeat_conf.fmap.lambda_": 1,
     "sds_conf.zoomout": 32,
     "diffusion.time": 1.0,
-    "opt.n_loop": 250,
+    "opt.n_loop": 200,
     "loss.sds": 1.0,
     "loss.proper": 1.0,
 }
@@ -144,13 +145,17 @@ class Datadicts:
     def __init__(self, shape_path, target_path):
         self.shape_path = shape_path
         basename_1 = os.path.basename(shape_path)
-        self.shape_dict, _ = helper.load_data(shape_path, "tmp/" + os.path.splitext(basename_1)[0]+".npz", "source", make_cache=True)
+        self.shape_dict, self.shape_dict_down = helper.load_data(shape_path, "tmp/" + os.path.splitext(basename_1)[0]+".npz", "source", make_cache=True)
         self.shape_surf = Surface(filename=shape_path)
+        self.shape_surf_down = Surface(filename=self.shape_dict_down["file"])
         self.target_path = target_path
         basename_2 = os.path.basename(target_path)
-        self.target_dict, _ = helper.load_data(target_path, "tmp/" + os.path.splitext(basename_2)[0]+".npz", "target", make_cache=True)
+        self.target_dict, self.target_dict_down = helper.load_data(target_path, "tmp/" + os.path.splitext(basename_2)[0]+".npz", "target", make_cache=True)
         self.target_surf = Surface(filename=target_path)
+        self.target_surf_down = Surface(filename=self.target_dict_down["file"])
+
         self.cmap1 = visu_pts(self.shape_surf)
+        self.cmap1_down = visu_pts(self.shape_surf_down)
 
 # -----------------------------
 # Gradio UI
@@ -189,16 +194,21 @@ def init_clicked(mesh1_path, mesh2_path,
         raise gr.Error("Please upload both meshes.")
     global datadicts
     datadicts = Datadicts(mesh1_path, mesh2_path)
-    shape_dict, target_dict = convert_dict(datadicts.shape_dict, 'cuda'), convert_dict(datadicts.target_dict, 'cuda')
+    shape_dict, target_dict = convert_dict(datadicts.shape_dict_down, 'cuda'), convert_dict(datadicts.target_dict_down, 'cuda')
     fmap_model_cuda = matcher.fmap_model.cuda()
     diff_model_cuda = matcher.diffusion_model
     diff_model_cuda.net.cuda()
     C12_pred_init, C21_pred_init, feat1, feat2, evecs_trans1, evecs_trans2 = fmap_model_cuda({"shape1": shape_dict, "shape2": target_dict}, diff_model=diff_model_cuda, scale=matcher.fmap_cfg.diffusion.time)
     C12_pred, C12_obj, mask_12 = C12_pred_init
-    p2p_init, _ = extract_p2p_torch_fmap(C12_obj, shape_dict["evecs"], target_dict["evecs"])
+    evecs1, evecs2 = torch.from_numpy(datadicts.shape_dict["evecs"]).cuda(), torch.from_numpy(datadicts.target_dict["evecs"]).cuda()
+    C_up, C_down = torch.from_numpy(datadicts.target_dict["Cup"]).cuda(), torch.from_numpy(datadicts.shape_dict_down["Cdown"]).cuda()
+    n_fmap = C12_obj.shape[-1]
+    with torch.no_grad():
+        C12_all = C_up.squeeze()[:n_fmap, :n_fmap] @ C12_obj.clone().squeeze() @ C_down.squeeze()[:n_fmap, :n_fmap]
+    p2p_init = FM_to_p2p(C12_all.cpu().numpy(), datadicts.shape_dict["evecs"], datadicts.target_dict["evecs"])
     return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_init, tag="init")
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=60)
 def run_clicked(mesh1_path, mesh2_path, yaml_path, lambda_val, zoomout_val, time_val, nloop_val, sds_val, proper_val, progress=gr.Progress(track_tqdm=True)):
     if not mesh1_path or not mesh2_path:
         raise gr.Error("Please upload both meshes.")
@@ -220,15 +230,19 @@ def run_clicked(mesh1_path, mesh2_path, yaml_path, lambda_val, zoomout_val, time
         if not (datadicts.shape_path == mesh1_path and datadicts.target_path == mesh2_path):
             datadicts = Datadicts(mesh1_path, mesh2_path)
 
-    shape_dict, target_dict = convert_dict(datadicts.shape_dict, 'cuda'), convert_dict(datadicts.target_dict, 'cuda')
-    target_normals = torch.from_numpy(datadicts.target_surf.surfel/np.linalg.norm(datadicts.target_surf.surfel, axis=-1, keepdims=True)).float().to("cuda")
+    shape_dict, target_dict = convert_dict(datadicts.shape_dict_down, 'cuda'), convert_dict(datadicts.target_dict_down, 'cuda')
+    target_normals = torch.from_numpy(datadicts.target_surf_down.surfel/np.linalg.norm(datadicts.target_surf_down.surfel, axis=-1, keepdims=True)).float().to("cuda")
 
     C12_new, p2p, p2p_init, _, loss_save = matcher.optimize(shape_dict, target_dict, target_normals)
-    evecs1, evecs2 = shape_dict["evecs"], target_dict["evecs"]
-    evecs_2trans = evecs2.t() @ torch.diag(target_dict["mass"])
+
+    C_up, C_down = torch.from_numpy(datadicts.target_dict["Cup"]).cuda(), torch.from_numpy(datadicts.shape_dict_down["Cdown"]).cuda()
+    evecs1, evecs2 = torch.from_numpy(datadicts.shape_dict["evecs"]).cuda(), torch.from_numpy(datadicts.target_dict["evecs"]).cuda()
+    evecs_2trans = evecs2.t() @ torch.diag(torch.from_numpy(datadicts.target_dict["mass"]).cuda())
     with torch.no_grad():
-        C12_end_zo = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_new.squeeze(), matcher.cfg.sds_conf.zoomout)
-    p2p_zo, _ = extract_p2p_torch_fmap(C12_end_zo, shape_dict["evecs"], target_dict["evecs"])
+        n_fmap = C12_new.shape[-1]
+        C12_all = C_up.squeeze()[:n_fmap, :n_fmap] @ C12_new.clone().squeeze() @ C_down.squeeze()[:n_fmap, :n_fmap]
+        C12_end_zo = torch_zoomout(evecs1, evecs2, evecs_2trans, C12_all, matcher.cfg.sds_conf.zoomout)
+        p2p_zo, _ = extract_p2p_torch_fmap(C12_end_zo, evecs1, evecs2)
     return build_outputs(datadicts.shape_surf, datadicts.target_surf, datadicts.cmap1, p2p_zo, tag="run")
 
 
@@ -327,5 +341,6 @@ if __name__ == "__main__":
     print("Matcher ready")
     #shutil.rmtree("tmp")
     os.makedirs("tmp", exist_ok=True)
+    os.makedirs("tmp/plys", exist_ok=True)
     datadicts = None
     demo.launch(share=args.share)

config/matching/sds.yaml

@@ -23,8 +23,8 @@ deepfeat_conf:
       abs: True
       normalize: False
       time: 1
-      batch_sds: 32
-      batch_mask: 200
+      batch_sds: 1
+      batch_mask: 16
 
 opt:
   n_loop: 300

notebook_helpers.py

@@ -1,27 +1,41 @@
 from utils.mesh import load_mesh
 from utils.geometry import get_operators, load_operators
 import os
-from utils.utils_func import convert_dict
+from utils.fmap import p2p_to_FM, KNNSearch
 from utils.surfaces import Surface
 import numpy as np
 
-def load_data(file, cache_path, name, num_evecs=128, make_cache=False, factor=None):
-    if factor is None:
-        verts_shape, faces, vnormals, area_shape, center_shape = load_mesh(file, return_vnormals=True)
-    else:
-        verts_shape, faces, vnormals, area_shape, center_shape = load_mesh(file, return_vnormals=True, scale=False)
-        verts_shape = verts_shape/factor
-        area_shape /= factor**2
+
+def get_down_p2ps(verts, verts_down):
+    tree = KNNSearch(verts)
+    ind_up = tree.query(verts_down, k=1).flatten()
+
+    tree_down = KNNSearch(verts_down)
+    ind_down = tree_down.query(verts, k=1).flatten()
+    return  ind_up, ind_down
+
+def load_data(file, cache_path, name, num_evecs=128, make_cache=False):
+    surf, surf_down, file_down = load_mesh(file)
+    cache_down = cache_path[:-4] + "_down" + ".npz"
     # print("Cache is: ", cache_path)
     if not os.path.exists(cache_path) or make_cache:
         print("Computing operators ...")
-        get_operators(verts_shape, faces, num_evecs, cache_path, vnormals)
+        vnormals = surf.computeVertexNormals()
+        get_operators(surf.vertices, surf.faces, num_evecs, cache_path, vnormals)
+        vnormals_down = surf_down.computeVertexNormals()
+        get_operators(surf_down.vertices, surf_down.faces, num_evecs, cache_down, vnormals_down)
     data_dict = load_operators(cache_path)
     data_dict['name'] = name
-    data_dict['normals'] = vnormals
-    data_dict['vertices'] = verts_shape
+    data_dict['vertices'] = surf.vertices
+    data_dict_down = load_operators(cache_down)
+    data_dict_down['name'] = name
+    data_dict_down['file'] = file_down
+    data_dict_down['vertices'] = surf_down.vertices
+    p2p_up, p2p_down = get_down_p2ps(surf.vertices, surf_down.vertices)
+    data_dict_down["Cdown"] = p2p_to_FM(p2p_up, data_dict['evecs'], data_dict_down['evecs'])
+    data_dict["Cup"] = p2p_to_FM(p2p_down, data_dict_down['evecs'], data_dict['evecs'])
     #batchify_dict(data_dict_torch)
-    return data_dict, area_shape
+    return data_dict, data_dict_down
 
 def get_map_info(file_1, file_2, dict_1, dict_2, dataset):
     shape_dict, target_dict = dict_1, dict_2

utils/fmap.py

@@ -4,13 +4,12 @@ import sys
 import numpy as np
 import scipy.linalg
 from tqdm import tqdm
+from scipy.spatial import cKDTree
 
 ROOT_DIR = osp.join(osp.abspath(osp.dirname(__file__)), '../')
 if ROOT_DIR not in sys.path:
     sys.path.append(ROOT_DIR)
 
-from utils_fmaps.misc import KNNSearch
-
 try:
     import pynndescent
     index = pynndescent.NNDescent(np.random.random((100, 3)), n_jobs=2)
@@ -19,6 +18,28 @@ try:
 except ImportError:
     ANN = False
 
+
+class KNNSearch(object):
+    DTYPE = np.float32
+    NJOBS = 4
+
+    def __init__(self, data):
+        self.data = np.asarray(data, dtype=self.DTYPE)
+        self.kdtree = cKDTree(self.data)
+
+    def query(self, kpts, k, return_dists=False):
+        kpts = np.asarray(kpts, dtype=self.DTYPE)
+        nndists, nnindices = self.kdtree.query(kpts, k=k, workers=self.NJOBS)
+        if return_dists:
+            return nnindices, nndists
+        else:
+            return nnindices
+
+    def query_ball(self, kpt, radius):
+        kpt = np.asarray(kpt, dtype=self.DTYPE)
+        assert kpt.ndim == 1
+        nnindices = self.kdtree.query_ball_point(kpt, radius, n_jobs=self.NJOBS)
+        return nnindices
 # https://github.com/RobinMagnet/pyFM
 
 

utils/geometry.py

@@ -477,7 +477,7 @@ def get_operators(verts, faces, k_eig=128, cache_path=None, normals=None, overwr
                 if (not np.array_equal(verts, cache_verts)) or (not np.array_equal(faces, cache_faces)):
                     i_cache_search += 1
                     print("hash collision! searching next.")
-                    continue
+                    overwrite_cache = True
 
                 # print("  cache hit!")
 
@@ -755,8 +755,8 @@ def get_operators_small(verts, faces, k_eig=128, cache_path=None, overwrite_cach
                 # If the cache doesn't match, keep looking
                 if (not np.array_equal(verts, cache_verts)) or (not np.array_equal(faces, cache_faces)):
                     i_cache_search += 1
-                    print("hash collision! searching next.")
-                    continue
+                    print("hash collision! overwriting.")
+                    overwrite_cache = True
 
                 # print("  cache hit!")
 

utils/mesh.py

@@ -10,7 +10,8 @@ import sys
 from utils.surfaces import Surface, centroid
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 from shape_data import get_data_dirs
-
+import vtk
+from vtk.util import numpy_support
 # List of file extensions to consider as "mesh" files.
 # Kudos to chatgpt!
 # Add or remove extensions here as needed.
@@ -104,24 +105,85 @@ def save_ply(file_name, V, F, Rho=None, color=None):
 
 
 
-def load_mesh(filepath, scale=True, return_vnormals=False):
+def load_mesh(filepath, scale=True):
     V, F = pp3d.read_mesh(filepath)
+    down_V, down_F = decimation(V, F)
+    down_surf = Surface(FV=[down_F,down_V])
+    down_file = os.path.join("tmp/plys", os.path.basename(filepath)[:-4])
+    save_ply(down_file, down_surf.vertices, down_surf.faces)
     surf = Surface(FV=[F,V])
     center, area = centroid(surf)
     if scale:
         new_vertices = (surf.vertices - center)/area
+        new_vertices_down = (down_surf.vertices - center) / area
     else:
         new_vertices = (surf.vertices - center)
-    vertices = np.asarray(new_vertices, dtype=np.float32)
-    faces = np.asarray(F, dtype=np.int32)
-
-    if return_vnormals:
-        vnormals = surf.computeVertexNormals()
-        if scale:
-            return vertices, faces, vnormals, area, center
-        return vertices, faces, vnormals, area, center
-    else:
-        return vertices, faces
+        new_vertices_down = (down_surf.vertices - center) / area
+    surf.updateVertices(np.array(new_vertices, dtype=np.float32))
+    down_surf.updateVertices(np.array(new_vertices_down, dtype=np.float32))
+    return surf, down_surf, down_file+".ply"
+
+
+def add_vectors(vertices: np.ndarray, faces: np.ndarray, vtk_vertices, vtk_faces):
+    assert len(vertices.shape) == 2
+    assert len(faces.shape) == 2
+    assert vertices.shape[1] == 3
+    assert faces.shape[1] == 3
+
+    # Add points
+    [num_vertices, _] = vertices.shape
+    for vertex_idx in range(num_vertices):
+        vtk_vertices.InsertNextPoint(vertices[vertex_idx, 0], vertices[vertex_idx, 1], vertices[vertex_idx, 2])
+    [num_faces, _] = faces.shape
+    for face_idx in range(num_faces):
+        vtk_faces.InsertNextCell(3)
+        for corner_idx in range(3):
+            vtk_faces.InsertCellPoint(faces[face_idx, corner_idx])
+    # Allocate additional memory
+    vtk_vertices.Resize(num_vertices)
+    # self._faces.Resize(self.num_faces + num_faces*3)
+
+    vtk_vertices.Modified()
+    vtk_faces.Modified()
+
+def decimation(verts, faces, target_n=1500):
+    vtk_vertices = vtk.vtkPoints()
+
+    # VTK polygone(surface) representation
+    vtk_faces = vtk.vtkCellArray()
+    surface_data = vtk.vtkPolyData()
+    surface_data.SetPoints(vtk_vertices)
+    surface_data.SetPolys(vtk_faces)
+    add_vectors(verts, faces, vtk_vertices, vtk_faces)
+
+
+    ratio = target_n /verts.shape[0]
+    print(ratio)
+    print("Before decimation\n"
+          "-----------------\n"
+          "There are " + str(surface_data.GetNumberOfPoints()) + "points.\n"
+          "There are " + str(surface_data.GetNumberOfPolys()) + "polygons.\n")
+
+    decimate = vtk.vtkQuadricDecimation()
+    decimate.SetInputData(surface_data)
+    decimate.SetTargetReduction(1-ratio)
+    decimate.Update()
+
+    decimatedPoly = vtk.vtkPolyData()
+    decimatedPoly.ShallowCopy(decimate.GetOutput())
+
+    print("After decimation \n"
+          "-----------------\n"
+          "There are " + str(decimatedPoly.GetNumberOfPoints()) + "points.\n"
+          "There are " + str(decimatedPoly.GetNumberOfPolys()) + "polygons.\n")
+    points = decimatedPoly.GetPoints().GetData()
+    vertices = numpy_support.vtk_to_numpy(points)  # shape: (n_points, 3)
+
+    # --- Faces (connectivity) ---
+    faces = decimatedPoly.GetPolys().GetData()
+    faces_np = numpy_support.vtk_to_numpy(faces)
+    faces_np = faces_np.reshape(-1, 4)[:, 1:]
+    return vertices, faces_np
 
 
 def mesh_geod_matrix(vertices, faces, do_tqdm=False, verbose=False):