Custom data - register known CAD model - how can I further optimize alignment?

[mrcfschr]

Hi there!
Thank you so much for this amazing work!
I'm trying to register known CAD models of anatomy to real world counterparts (patients) based solely on the partially occluded and noisy reconstructed point cloud data I get from the ToF camera of the HoloLens 2 mixed reality headset(using no RGB - only D). I'm completely new to Deep Learning, and I was wondering how I can best leverage the fact that I have the mesh of the 3D model I'm trying to register beforehand. Ideally I want to create a pipeline that lets one upload a 3D model and do a bunch of processing before to create the best possible network to register the CAD model.
Before doing any retraining, is there anything I can change in my config file to optimize the pretrained network for this use case of aligning smaller sized objects like the CT scan of a knee? I used the indoor config in my test and only changed dgcnn_k to 4 to avoid an error.(see What I have done so far ) Do you have advise what else I could optimize/ change? Unfortunately out of the box the network occasionally registers the CAD model upside down.

Would it be a good idea to train this network based solely on the one CAD model that I'm trying to register? Should I collect a lot of noise data that one would encounter in my use case and then train a network to be able to generate realistic noise that can be used during this custom training step?
Thank you so much for your consideration and sharing your work!

What I have done so far (for any complete beginners wondering how to use this for custom data):

SetUP before demo.py

• on Ubuntu 20.04.3 LTS with the right devtools installed git, ninja etc
• git clone https://github.com/overlappredator/OverlapPredator.git
• conda activate py3.8.5 (environment with all requirements installed)
• cd OverlapPredator
• cd cpp_wrappers
• sh compile_wrappers.sh
• cd ..
• sh scripts/download_data_weight.sh
• python scripts/demo.py configs/test/indoor.yaml

demo.py works like a charm!.

Using custom data:

• I want to align 3D model (source.ply) of my arm with a sensor reading (Target.ply) of it
• files can be found here: https://drive.google.com/drive/folders/1sbIwYgY4YODjrLvshC6-q_010YuCv9DU?usp=sharing
• download and place into the asset folder
• comment out lines 47 & 48 in scripts/demo.py
• uncomment lines 51 -56 in scripts/demo.py like so:

def __getitem__(self,item): 
       # get pointcloud
       # src_pcd = torch.load(self.src_path).astype(np.float32)
       # tgt_pcd = torch.load(self.tgt_path).astype(np.float32)   
       
       
       src_pcd = o3d.io.read_point_cloud(self.src_path)
       tgt_pcd = o3d.io.read_point_cloud(self.tgt_path)
       src_pcd = src_pcd.voxel_down_sample(0.025)
       tgt_pcd = tgt_pcd.voxel_down_sample(0.025)
       src_pcd = np.array(src_pcd.points).astype(np.float32)
       tgt_pcd = np.array(tgt_pcd.points).astype(np.float32)


       src_feats=np.ones_like(src_pcd[:,:1]).astype(np.float32)
       tgt_feats=np.ones_like(tgt_pcd[:,:1]).astype(np.float32)

       # fake the ground truth information
       rot = np.eye(3).astype(np.float32)
       trans = np.ones((3,1)).astype(np.float32)
       correspondences = torch.ones(1,2).long()

       return src_pcd,tgt_pcd,src_feats,tgt_feats,rot,trans, correspondences, src_pcd, tgt_pcd, torch.ones(1)

• copy indoor.yaml and rename to testconfig.py
• change line 73 & 74 to point to source and target point cloud like so:

demo:
  src_pcd: assets/Source.ply
  tgt_pcd: assets/Target.ply
  n_points: 1000

• run python scripts/demo.py configs/test/testconfig.yaml
• cause "RuntimeError: invalid argument 5: k not in range for dimension at /pytorch/aten/src/THC/generic/THCTensorTopK.cu:26"
• try setting dgcnn_k from 10 to a lower number (4 worked for me) in testconfig.yaml like so:

overlap_attention_module:
  gnn_feats_dim: 256 
  dgcnn_k: 4
  num_head: 4
  nets: ['self','cross','self']

Results:

[ShengyuH]

hi,

Thanks for your interest and encouraging words.

To answer your question:

• Due to the symmetry in your data, "registering objects upside down" is quite expected. If you check those deep models registering ModelNet40 datasets, quite often, they restrict the test samples to have rotation angles within 45 degrees such that they can directly avoid ambiguity caused by symmetry. In large scale data, we rely on larger context to disambiguate them, I guess this would be challenging in your case.

• Without some hacks, KPConv backbone is sensitive to hyper-parameter changes, so not much you can do with the pre-trained models. However, you can check D3Feat paper to see how they hack the backbone to achieve generalisation, though only to some extent.

• I think finetuning the pretrained model on your collected data is a good idea.

Let me know if your need more assistance. I was once asked about such application scenarios and I am looking forward to see how it works.

best,
Shengyu