prs-eth / OverlapPredator

[CVPR 2021, Oral] PREDATOR: Registration of 3D Point Clouds with Low Overlap.
https://shengyuh.github.io/predator/index.html
MIT License
505 stars 70 forks source link

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

Closed mrcfschr closed 3 years ago

mrcfschr commented 3 years ago

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: Screenshot from 2021-09-06 15-57-48 Screenshot from 2021-09-06 15-57-52 Screenshot from 2021-09-06 15-57-45

ShengyuH commented 3 years ago

hi,

Thanks for your interest and encouraging words.

To answer your question:

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