├── Framework.png ├── README.md ├── __init__.py ├── __pycache__ ├── centerloss.cpython-36.pyc ├── centerloss.cpython-37.pyc ├── centerloss.cpython-38.pyc ├── centerloss_old.cpython-36.pyc ├── centerloss_old.cpython-37.pyc ├── nt_xent.cpython-36.pyc ├── nt_xent.cpython-37.pyc └── nt_xent.cpython-38.pyc ├── baseline_pc_train.py ├── baseline_train.py ├── centerloss.py ├── data_utils ├── .ipynb_checkpoints │ └── ModelNetDataLoader-checkpoint.py ├── ModelNetDataLoader.py ├── S3DISDataLoader.py ├── ShapeNetDataLoader.py ├── ShapeNetDataLoader.pyc ├── __init__.py ├── __init__.pyc ├── __pycache__ │ ├── S3DISDataLoader.cpython-37.pyc │ ├── ShapeNetDataLoader.cpython-37.pyc │ ├── __init__.cpython-37.pyc │ └── indoor3d_util.cpython-37.pyc ├── collect_indoor3d_data.py ├── indoor3d_util.py └── meta │ ├── anno_paths.txt │ └── class_names.txt ├── feature_extractor.py ├── feature_vis.py ├── finetune_pointnet.py ├── finetune_pointnet_freeze.py ├── models ├── .ipynb_checkpoints │ ├── SVCNN-checkpoint.py │ ├── SVCNN_old-checkpoint.py │ ├── SVCNN_twfview-checkpoint.py │ ├── dgcnn-checkpoint.py │ ├── meshnet-checkpoint.py │ ├── pointnet-checkpoint.py │ ├── pointnet_fc-checkpoint.py │ └── resnet-checkpoint.py ├── SVCNN-Copy1.py ├── SVCNN.py ├── SVCNN_old.py ├── SVCNN_twfview.py ├── __init__.py ├── __init__.pyc ├── __pycache__ │ ├── SVCNN.cpython-36.pyc │ ├── SVCNN.cpython-37.pyc │ ├── SVCNN.cpython-38.pyc │ ├── SVCNN_twfview.cpython-36.pyc │ ├── SVCNN_twfview.cpython-37.pyc │ ├── SVCNN_twfview.cpython-38.pyc │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── dgcnn.cpython-36.pyc │ ├── dgcnn.cpython-37.pyc │ ├── dgcnn.cpython-38.pyc │ ├── meshnet.cpython-36.pyc │ ├── meshnet.cpython-37.pyc │ ├── meshnet.cpython-38.pyc │ ├── pointnet.cpython-36.pyc │ ├── pointnet.cpython-37.pyc │ ├── pointnet.cpython-38.pyc │ ├── pointnet_fc.cpython-37.pyc │ ├── pointnet_part_seg.cpython-36.pyc │ ├── pointnet_part_seg.cpython-37.pyc │ ├── pointnet_part_seg.cpython-38.pyc │ ├── resnet.cpython-36.pyc │ ├── resnet.cpython-37.pyc │ └── resnet.cpython-38.pyc ├── dgcnn-Copy1.py ├── dgcnn.py ├── dgcnn.pyc ├── meshnet-Copy1.py ├── meshnet.py ├── pointnet.py ├── pointnet_fc.py ├── pointnet_part_seg.py ├── pointnet_util.py └── resnet.py ├── nt_xent.py ├── pc_test.py ├── pc_train.py ├── pc_train_centerloss.py ├── provider.py ├── retrieve_eval.py ├── test.py ├── test_12views.py ├── test_distance.py ├── test_point.py ├── tools ├── .ipynb_checkpoints │ ├── dual_dataloader-checkpoint.py │ ├── test_dataloader-checkpoint.py │ ├── test_dataloader_twfview-checkpoint.py │ ├── triplet_dataloader-checkpoint.py │ ├── triplet_dataloader_twfview-checkpoint.py │ └── visualize-checkpoint.py ├── .python-version ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-36.pyc │ ├── __init__.cpython-37.pyc │ ├── __init__.cpython-38.pyc │ ├── dual_dataloader.cpython-36.pyc │ ├── dual_dataloader.cpython-37.pyc │ ├── dual_dataloader.cpython-38.pyc │ ├── test_dataloader.cpython-36.pyc │ ├── test_dataloader.cpython-37.pyc │ ├── test_dataloader.cpython-38.pyc │ ├── test_dataloader_twfview.cpython-36.pyc │ ├── test_dataloader_twfview.cpython-37.pyc │ ├── test_dataloader_twfview.cpython-38.pyc │ ├── triplet_dataloader.cpython-36.pyc │ ├── triplet_dataloader.cpython-37.pyc │ ├── triplet_dataloader.cpython-38.pyc │ ├── triplet_dataloader_twfview.cpython-36.pyc │ ├── triplet_dataloader_twfview.cpython-37.pyc │ ├── utils.cpython-36.pyc │ ├── utils.cpython-37.pyc │ ├── utils.cpython-38.pyc │ ├── visualize.cpython-36.pyc │ ├── visualize.cpython-37.pyc │ └── visualize.cpython-38.pyc ├── build.sh ├── dual_dataloader.py ├── gen_cloud_images.py ├── gen_cloud_images_shapenet.py ├── render_balls_so.cpp ├── render_balls_so.so ├── test_dataloader.py ├── test_dataloader_twfview.py ├── triplet_dataloader.py ├── triplet_dataloader_twfview.py ├── utils.py └── visualize.py ├── train.py ├── tsne.py └── wandb └── settings /Framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/Framework.png -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Multimodal Semi-Supervised Learning for 3D Objects 2 | 3 | Our paper has been accepted by BMVC 2021 4 | 5 | Arvix version [here](https://arxiv.org/abs/2110.11601) 6 | 7 | ![image](https://github.com/AutoAILab/Multimodal-Semi-Learning/blob/main/Framework.png) 8 | 9 | ## Abstract 10 | 11 | We propose a novel multimodal semi-supervised learning framework by introducing instance-level consistency constraint and a novel multimodal contrastive prototype (M2CP) loss. The instance-level consistency enforces the network to generate consistent representations for multimodal data of the same object regardless of its modality. The M2CP maintains a multimodal prototype for each class and learns features with small intra-class variations by minimizing the feature distance of each object to its prototype while maximizing the distance to the others. Our proposed framework significantly outperforms all the state-of-the-art counterparts for both classification and retrieval tasks by a large margin on the modelNet10 and ModelNet40 datasets. 12 | 13 | ## Download Dataset 14 | [Image](https://drive.google.com/file/d/1-Gz4HMwM96zEnR3Hji4UEAGB4FPrd5Xo/view?usp=sharing) 15 | 16 | [Mesh](https://drive.google.com/file/d/1pcOK10d4NrgJzAbqLJFF9hIA43pwiVEH/view?usp=sharing) 17 | 18 | [Point Clouds](https://drive.google.com/file/d/1VyXh4Q7EJK0uUM6xAgAHyzwpFFui_se4/view?usp=sharing) 19 | 20 | 21 | ## Installation 22 | Install ```Python``` -- This repo is tested with Python 3.7.6. 23 | 24 | Install ```NumPy``` -- This repo is tested with NumPy 1.18.5. Please make sure your NumPy version is at least 1.18. 25 | 26 | Install ```PyTorch``` with CUDA -- This repo is tested with PyTorch 1.5.1, CUDA 10.2. It may work with newer versions, but that is not guaranteed. A lower version may be problematic. 27 | 28 | Install ```TensorFlow``` (for TensorBoard) -- This repo is tested with TensorFlow 2.2.0. 29 | 30 | 31 | ## Training 32 | This netowrk is trained with two 16G Tesla V100 GPU 33 | 34 | Remember change the dataloader root before training 35 | ``` 36 | python train.py 37 | ``` 38 | 39 | ## Testing 40 | The trained model is save in checkpoints file. Remember change the test model name in ```test_12views.py``` before testing. 41 | ``` 42 | python test_12views.py 43 | ``` 44 | ## Pretrained model 45 | We provide the pre-trained models of [ModelNet40-10%](https://drive.google.com/file/d/1StclCQu9ePoFPGh8jFQWRkYQkYyUxm_9/view?usp=share_link) 46 | 47 | ## Citation 48 | ``` 49 | @article{chen2021multimodal, 50 | title={Multimodal Semi-Supervised Learning for 3D Objects}, 51 | author={Chen, Zhimin and Jing, Longlong and Liang, Yang and Tian, YingLi and Li, Bing}, 52 | journal={arXiv preprint arXiv:2110.11601}, 53 | year={2021} 54 | } 55 | ``` 56 | 57 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__init__.py -------------------------------------------------------------------------------- /__pycache__/centerloss.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/centerloss.cpython-36.pyc -------------------------------------------------------------------------------- /__pycache__/centerloss.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/centerloss.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/centerloss.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/centerloss.cpython-38.pyc -------------------------------------------------------------------------------- /__pycache__/centerloss_old.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/centerloss_old.cpython-36.pyc -------------------------------------------------------------------------------- /__pycache__/centerloss_old.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/centerloss_old.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/nt_xent.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/nt_xent.cpython-36.pyc -------------------------------------------------------------------------------- /__pycache__/nt_xent.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/nt_xent.cpython-37.pyc -------------------------------------------------------------------------------- /__pycache__/nt_xent.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/__pycache__/nt_xent.cpython-38.pyc -------------------------------------------------------------------------------- /baseline_pc_train.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import 2 | from models.dgcnn import DGCNN 3 | from models.pointnet_part_seg import PointnetPartSeg 4 | # from models.pointnet_part_seg import PointNet_Part 5 | from models.meshnet import MeshNet 6 | from models.SVCNN import Semi3D, SingleViewNet 7 | from tools.triplet_dataloader import TripletDataloader 8 | from tools.utils import calculate_accuracy 9 | import numpy as np 10 | import math 11 | import os 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | from torch.autograd import Variable 16 | import torchvision.models as models 17 | import argparse 18 | import torch.optim as optim 19 | import time 20 | 21 | from nt_xent import NTXentLoss 22 | 23 | from torch.utils.tensorboard import SummaryWriter 24 | 25 | import warnings 26 | warnings.filterwarnings('ignore',category=FutureWarning) 27 | 28 | 29 | 30 | def training(args): 31 | if not os.path.exists(args.save): 32 | os.makedirs(args.save) 33 | 34 | # img_net = SingleViewNet(pre_trained = True) 35 | pt_net = DGCNN(args) 36 | # pt_net = PointnetPartSeg() 37 | 38 | # meshnet = MeshNet() 39 | 40 | # model = Semi3D(img_net, pt_net, meshnet) 41 | 42 | pt_net = pt_net.to('cuda') 43 | 44 | pt_net = torch.nn.DataParallel(pt_net) 45 | 46 | pt_net.train(True) 47 | 48 | optimizer = optim.SGD(pt_net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) 49 | 50 | writer = SummaryWriter(os.path.join(args.save, 'summary')) 51 | 52 | #data splittted into unlabeled/labeled/test 53 | labeled_set = TripletDataloader(dataset = 'ModelNet40', num_points = args.num_points, partition='labeled', perceptange = 5) 54 | labeled_data_loader = torch.utils.data.DataLoader(labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=8, drop_last=True) 55 | 56 | 57 | # unlabeled_set = TripletDataloader(dataset = 'ModelNet40', num_points = args.num_points, partition='unlabeled', perceptange = 5) 58 | # unlabeled_data_loader = torch.utils.data.DataLoader(unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=8, drop_last=True) 59 | 60 | print('************************************************************') 61 | print(' check the following important parametes ') 62 | print('the number of labeled sample: ', len(labeled_set)) 63 | # print('the number of unlabeled sample: ', len(unlabeled_set)) 64 | print('the temperature for the probability: ', args.T) 65 | print('the threshold for the probability: ', args.threshold) 66 | print('************************************************************') 67 | 68 | # The loss introduced in Hinton's paper 69 | # nt_xent_criterion = NTXentLoss('cuda', args.batch_size, temperature = 0.5, use_cosine_similarity = True) 70 | # mse_criterion = nn.MSELoss() 71 | ce_criterion = nn.CrossEntropyLoss(reduction='mean') 72 | 73 | iteration = 0 74 | start_time = time.time() 75 | for epoch in range(args.epochs): 76 | for l_data in labeled_data_loader: 77 | pt1, img1, img1V, centers1, corners1, normals1, neighbor_index1, target1 = l_data #the last one is the target 78 | 79 | pt1 = Variable(pt1).to('cuda') 80 | pt1 = pt1.permute(0,2,1) 81 | 82 | target1 = torch.squeeze(target1) 83 | target1 = Variable(target1).to('cuda') 84 | 85 | optimizer.zero_grad() 86 | # 87 | # pt_pred1, mesh_pred1, img_pred1, fused_pred1, pt_feat1, mesh_feat1, img_feat1 = model(pt1, img1, img1V, centers1, corners1, normals1, neighbor_index1) 88 | pt_pred1, pt_feat1, pt_base1 = pt_net(pt1) 89 | # pt_pred2, mesh_pred2, img_pred2, fused_pred2, pt_feat2, mesh_feat2, img_feat2 = model(pt2, img2, img2V, centers2, corners2, normals2, neighbor_index2) 90 | 91 | #cross-entropy loss on the labeled data 92 | pt_ce_loss = ce_criterion(pt_pred1, target1) 93 | loss = pt_ce_loss 94 | 95 | loss.backward() 96 | 97 | #update the parameters for the center_loss 98 | optimizer.step() 99 | 100 | pt_acc1 = calculate_accuracy(pt_pred1, target1) 101 | 102 | #classification accuracy on the labeld sample 103 | writer.add_scalar('pt_acc', pt_acc1, iteration) 104 | 105 | #Xentropy loss on the labeled data 106 | writer.add_scalar('pt_ce_loss', pt_ce_loss.item(), iteration) 107 | 108 | writer.add_scalar('loss/loss', loss.item(), iteration) 109 | 110 | 111 | if (iteration%args.lr_step) == 0: 112 | lr = args.lr * (0.1 ** (iteration // args.lr_step)) 113 | print('New LR: ' + str(lr)) 114 | for param_group in optimizer.param_groups: 115 | param_group['lr'] = lr 116 | 117 | if iteration % args.per_print == 0: 118 | print('[%d][%d] loss: %.2f acc: %.2f time: %.2f vid: %d' % \ 119 | (epoch, iteration, loss.item(), pt_acc1, time.time() - start_time, 2 * pt1.size(0))) 120 | start_time = time.time() 121 | 122 | iteration = iteration + 1 123 | if((iteration+1) % args.per_save) ==0: 124 | print('----------------- Save The Network ------------------------') 125 | 126 | pt_net_name = args.save + str(iteration+1)+'-pt_net.pkl' 127 | torch.save({'state_dict': pt_net.state_dict()}, pt_net_name) 128 | 129 | iteration = iteration + 1 130 | if iteration > args.max_step: 131 | return 132 | 133 | 134 | if __name__ == "__main__": 135 | # Training settings 136 | parser = argparse.ArgumentParser(description='Learning View and Model invariant features for 3D shapes') 137 | 138 | parser.add_argument('--batch_size', type=int, default=48, metavar='batch_size', 139 | help='Size of batch)') 140 | 141 | parser.add_argument('--epochs', type=int, default=2000, metavar='N', 142 | help='number of episode to train ') 143 | #optimizer 144 | parser.add_argument('--lr', type=float, default=0.005, metavar='LR', 145 | help='learning rate (default: 0.001, 0.1 if using sgd)') 146 | 147 | parser.add_argument('--lr_step', type=int, default = 2000, 148 | help='how many iterations to decrease the learning rate') 149 | 150 | parser.add_argument('--max_step', type=int, default = 6100, 151 | help='maximum steps to train the network') 152 | 153 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 154 | help='SGD momentum (default: 0.9)') 155 | 156 | parser.add_argument('--no_cuda', type=bool, default=False, 157 | help='enables CUDA training') 158 | 159 | parser.add_argument('--seed', type=int, default=1, metavar='S', 160 | help='random seed (default: 1)') 161 | 162 | parser.add_argument('--T', type=int, default = 1, 163 | help='temperature for the prediction') 164 | 165 | parser.add_argument('--threshold', type=int, default = 0.95, 166 | help='threshold for the positive samples') 167 | 168 | #image for SVCNN 169 | parser.add_argument('--num_views', type=int, default=180, metavar='S', 170 | help='number of views for training (default: 6)') 171 | #DGCNN 172 | parser.add_argument('--num_points', type=int, default=1024, 173 | help='num of points to use') 174 | parser.add_argument('--dropout', type=float, default=0.5, 175 | help='dropout rate') 176 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 177 | help='Dimension of embeddings') 178 | parser.add_argument('--k', type=int, default=20, metavar='N', 179 | help='Num of nearest neighbors to use') 180 | 181 | parser.add_argument('--model_path', type=str, default='', metavar='N', 182 | help='Pretrained model path') 183 | 184 | parser.add_argument('--weight_decay', type=float, default=1e-3, metavar='weight_decay', 185 | help='learning rate (default: 1e-3)') 186 | 187 | parser.add_argument('--per_save', type=int, default=2000, 188 | help='how many iterations to save the model') 189 | 190 | parser.add_argument('--per_print', type=int, default=100, 191 | help='how many iterations to print the loss and accuracy') 192 | parser.add_argument('--save', type=str, default='./checkpoints/ModelNet40-pt1024-5percent-supervised/', 193 | help='path to save the final model') 194 | 195 | parser.add_argument('--gpu_id', type=str, default='0,1,2', 196 | help='GPU used to train the network') 197 | 198 | parser.add_argument('--log', type=str, default='log/', 199 | help='path to the log information') 200 | 201 | args = parser.parse_args() 202 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 203 | torch.backends.cudnn.enabled = False 204 | training(args) 205 | -------------------------------------------------------------------------------- /centerloss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | class CenterLoss(nn.Module): 6 | """Center loss. 7 | 8 | Reference: 9 | Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016. 10 | 11 | Args: 12 | num_classes (int): number of classes. 13 | feat_dim (int): feature dimension. 14 | """ 15 | def __init__(self, num_classes = 40, feat_dim = 256,temperature = 0.07, use_gpu = True): 16 | super(CenterLoss, self).__init__() 17 | self.num_classes = num_classes 18 | self.feat_dim = feat_dim 19 | self.use_gpu = use_gpu 20 | self.temperature=temperature 21 | self.base_temperature=0.07 22 | if self.use_gpu: 23 | self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda()) 24 | else: 25 | self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim)) 26 | 27 | def forward(self, x, labels, valid_mask): 28 | """ 29 | Args: 30 | x: feature matrix with shape (batch_size, feat_dim). 31 | labels: ground truth labels with shape (batch_size). 32 | """ 33 | batch_size = x.size(0) 34 | 35 | x = F.normalize(x, dim=1) 36 | centers = F.normalize(self.centers, dim=1) 37 | 38 | anchor_dot_contrast = torch.div( 39 | torch.matmul(x, centers.T), 40 | self.temperature) 41 | 42 | 43 | # x = F.normalize(x, dim=1) 44 | # centers = F.normalize(self.centers, dim=1) 45 | 46 | # anchor_dot_contrast = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \ 47 | # torch.pow(centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t() 48 | # anchor_dot_contrast.addmm_(1, -2, x, centers.t()) 49 | 50 | # anchor_dot_contrast = torch.div(anchor_dot_contrast, self.temperature) 51 | 52 | logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True) 53 | logits = anchor_dot_contrast - logits_max.detach() 54 | 55 | classes = torch.arange(self.num_classes).long() 56 | if self.use_gpu: classes = classes.cuda() 57 | labels = labels.unsqueeze(1).expand(batch_size, self.num_classes) 58 | mask = labels.eq(classes.expand(batch_size, self.num_classes)) 59 | mask = mask.float() 60 | valid_mask = valid_mask.float() 61 | exp_logits = torch.exp(logits) 62 | log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True)) 63 | mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1) 64 | 65 | # loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos 66 | loss = - (0.5/ self.base_temperature) * mean_log_prob_pos 67 | 68 | loss = loss * valid_mask 69 | loss = loss.mean() 70 | 71 | return loss, centers 72 | 73 | 74 | 75 | 76 | # import torch 77 | # import torch.nn as nn 78 | 79 | # class CenterLoss(nn.Module): 80 | # """Center loss. 81 | # 82 | # Reference: 83 | # Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016. 84 | # 85 | # Args: 86 | # num_classes (int): number of classes. 87 | # feat_dim (int): feature dimension. 88 | # """ 89 | # def __init__(self, num_classes=10, feat_dim=2, use_gpu=True): 90 | # super(CenterLoss, self).__init__() 91 | # self.num_classes = num_classes 92 | # self.feat_dim = feat_dim 93 | # self.use_gpu = use_gpu 94 | # 95 | # if self.use_gpu: 96 | # self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda()) 97 | # else: 98 | # self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim)) 99 | # 100 | # def forward(self, x, labels): 101 | # """ 102 | # Args: 103 | # x: feature matrix with shape (batch_size, feat_dim). 104 | # labels: ground truth labels with shape (batch_size). 105 | # """ 106 | # batch_size = x.size(0) 107 | # distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \ 108 | # torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t() 109 | # distmat.addmm_(1, -2, x, self.centers.t()) 110 | # 111 | # classes = torch.arange(self.num_classes).long() 112 | # if self.use_gpu: classes = classes.cuda() 113 | # labels = labels.unsqueeze(1).expand(batch_size, self.num_classes) 114 | # mask = labels.eq(classes.expand(batch_size, self.num_classes)) 115 | # 116 | # dist = distmat * mask.float() 117 | # loss = dist.clamp(min=1e-12, max=1e+12).sum() / batch_size 118 | # 119 | # return loss -------------------------------------------------------------------------------- /data_utils/.ipynb_checkpoints/ModelNetDataLoader-checkpoint.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import warnings 3 | import os 4 | from torch.utils.data import Dataset 5 | warnings.filterwarnings('ignore') 6 | 7 | 8 | 9 | def pc_normalize(pc): 10 | centroid = np.mean(pc, axis=0) 11 | pc = pc - centroid 12 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 13 | pc = pc / m 14 | return pc 15 | 16 | def farthest_point_sample(point, npoint): 17 | """ 18 | Input: 19 | xyz: pointcloud data, [N, D] 20 | npoint: number of samples 21 | Return: 22 | centroids: sampled pointcloud index, [npoint, D] 23 | """ 24 | N, D = point.shape 25 | xyz = point[:,:3] 26 | centroids = np.zeros((npoint,)) 27 | distance = np.ones((N,)) * 1e10 28 | farthest = np.random.randint(0, N) 29 | for i in range(npoint): 30 | centroids[i] = farthest 31 | centroid = xyz[farthest, :] 32 | dist = np.sum((xyz - centroid) ** 2, -1) 33 | mask = dist < distance 34 | distance[mask] = dist[mask] 35 | farthest = np.argmax(distance, -1) 36 | point = point[centroids.astype(np.int32)] 37 | return point 38 | 39 | class ModelNetDataLoader(Dataset): 40 | def __init__(self, root, npoint=1024, split='train', uniform=False, normal_channel=True, cache_size=15000): 41 | self.root = root 42 | self.npoints = npoint 43 | self.uniform = uniform 44 | self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt') 45 | 46 | self.cat = [line.rstrip() for line in open(self.catfile)] 47 | self.classes = dict(zip(self.cat, range(len(self.cat)))) 48 | self.normal_channel = normal_channel 49 | 50 | shape_ids = {} 51 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))] 52 | shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))] 53 | 54 | assert (split == 'train' or split == 'test') 55 | shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]] 56 | # list of (shape_name, shape_txt_file_path) tuple 57 | self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i 58 | in range(len(shape_ids[split]))] 59 | print('The size of %s data is %d'%(split,len(self.datapath))) 60 | 61 | self.cache_size = cache_size # how many data points to cache in memory 62 | self.cache = {} # from index to (point_set, cls) tuple 63 | 64 | def __len__(self): 65 | return len(self.datapath) 66 | 67 | def _get_item(self, index): 68 | if index in self.cache: 69 | point_set, cls = self.cache[index] 70 | else: 71 | fn = self.datapath[index] 72 | cls = self.classes[self.datapath[index][0]] 73 | cls = np.array([cls]).astype(np.int32) 74 | point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32) 75 | if self.uniform: 76 | point_set = farthest_point_sample(point_set, self.npoints) 77 | else: 78 | point_set = point_set[0:self.npoints,:] 79 | 80 | point_set[:, 0:3] = pc_normalize(point_set[:, 0:3]) 81 | 82 | if not self.normal_channel: 83 | point_set = point_set[:, 0:3] 84 | 85 | if len(self.cache) < self.cache_size: 86 | self.cache[index] = (point_set, cls) 87 | 88 | return point_set, cls 89 | 90 | def __getitem__(self, index): 91 | return self._get_item(index) 92 | 93 | 94 | 95 | 96 | if __name__ == '__main__': 97 | import torch 98 | 99 | data = ModelNetDataLoader('/data/modelnet40_normal_resampled/',split='train', uniform=False, normal_channel=True,) 100 | DataLoader = torch.utils.data.DataLoader(data, batch_size=12, shuffle=True) 101 | for point,label in DataLoader: 102 | print(point.shape) 103 | print(label.shape) -------------------------------------------------------------------------------- /data_utils/ModelNetDataLoader.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import warnings 3 | import os 4 | from torch.utils.data import Dataset 5 | warnings.filterwarnings('ignore') 6 | 7 | 8 | 9 | def pc_normalize(pc): 10 | centroid = np.mean(pc, axis=0) 11 | pc = pc - centroid 12 | m = np.max(np.sqrt(np.sum(pc**2, axis=1))) 13 | pc = pc / m 14 | return pc 15 | 16 | def farthest_point_sample(point, npoint): 17 | """ 18 | Input: 19 | xyz: pointcloud data, [N, D] 20 | npoint: number of samples 21 | Return: 22 | centroids: sampled pointcloud index, [npoint, D] 23 | """ 24 | N, D = point.shape 25 | xyz = point[:,:3] 26 | centroids = np.zeros((npoint,)) 27 | distance = np.ones((N,)) * 1e10 28 | farthest = np.random.randint(0, N) 29 | for i in range(npoint): 30 | centroids[i] = farthest 31 | centroid = xyz[farthest, :] 32 | dist = np.sum((xyz - centroid) ** 2, -1) 33 | mask = dist < distance 34 | distance[mask] = dist[mask] 35 | farthest = np.argmax(distance, -1) 36 | point = point[centroids.astype(np.int32)] 37 | return point 38 | 39 | class ModelNetDataLoader(Dataset): 40 | def __init__(self, root, npoint=1024, split='train', uniform=False, normal_channel=True, cache_size=15000): 41 | self.root = root 42 | self.npoints = npoint 43 | self.uniform = uniform 44 | self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt') 45 | 46 | self.cat = [line.rstrip() for line in open(self.catfile)] 47 | self.classes = dict(zip(self.cat, range(len(self.cat)))) 48 | self.normal_channel = normal_channel 49 | 50 | shape_ids = {} 51 | shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_train.txt'))] 52 | shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 'modelnet40_test.txt'))] 53 | 54 | assert (split == 'train' or split == 'test') 55 | shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]] 56 | # list of (shape_name, shape_txt_file_path) tuple 57 | self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], shape_ids[split][i]) + '.txt') for i 58 | in range(len(shape_ids[split]))] 59 | print('The size of %s data is %d'%(split,len(self.datapath))) 60 | 61 | self.cache_size = cache_size # how many data points to cache in memory 62 | self.cache = {} # from index to (point_set, cls) tuple 63 | 64 | def __len__(self): 65 | return len(self.datapath) 66 | 67 | def _get_item(self, index): 68 | if index in self.cache: 69 | point_set, cls = self.cache[index] 70 | else: 71 | fn = self.datapath[index] 72 | cls = self.classes[self.datapath[index][0]] 73 | cls = np.array([cls]).astype(np.int32) 74 | point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32) 75 | if self.uniform: 76 | point_set = farthest_point_sample(point_set, self.npoints) 77 | else: 78 | point_set = point_set[0:self.npoints,:] 79 | 80 | point_set[:, 0:3] = pc_normalize(point_set[:, 0:3]) 81 | 82 | if not self.normal_channel: 83 | point_set = point_set[:, 0:3] 84 | 85 | if len(self.cache) < self.cache_size: 86 | self.cache[index] = (point_set, cls) 87 | 88 | return point_set, cls 89 | 90 | def __getitem__(self, index): 91 | return self._get_item(index) 92 | 93 | 94 | 95 | 96 | if __name__ == '__main__': 97 | import torch 98 | 99 | data = ModelNetDataLoader('/data/modelnet40_normal_resampled/',split='train', uniform=False, normal_channel=True,) 100 | DataLoader = torch.utils.data.DataLoader(data, batch_size=12, shuffle=True) 101 | for point,label in DataLoader: 102 | print(point.shape) 103 | print(label.shape) -------------------------------------------------------------------------------- /data_utils/ShapeNetDataLoader.py: -------------------------------------------------------------------------------- 1 | # *_*coding:utf-8 *_* 2 | import os 3 | import json 4 | import warnings 5 | import numpy as np 6 | from torch.utils.data import Dataset 7 | import random 8 | 9 | warnings.filterwarnings('ignore') 10 | 11 | def pc_normalize(pc): 12 | centroid = np.mean(pc, axis=0) 13 | pc = pc - centroid 14 | m = np.max(np.sqrt(np.sum(pc ** 2, axis=1))) 15 | pc = pc / m 16 | return pc 17 | 18 | class PartNormalDataset(Dataset): 19 | def __init__(self,root = './dataset/shapenetcore_partanno_segmentation_benchmark_v0_normal', npoints=2500, split='train', class_choice=None, normal_channel=False): 20 | self.npoints = npoints 21 | self.root = root 22 | self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') 23 | self.cat = {} 24 | self.normal_channel = normal_channel 25 | 26 | 27 | with open(self.catfile, 'r') as f: 28 | for line in f: 29 | ls = line.strip().split() 30 | self.cat[ls[0]] = ls[1] 31 | self.cat = {k: v for k, v in self.cat.items()} 32 | self.classes_original = dict(zip(self.cat, range(len(self.cat)))) 33 | 34 | if not class_choice is None: 35 | self.cat = {k:v for k,v in self.cat.items() if k in class_choice} 36 | # print(self.cat) 37 | 38 | self.meta = {} 39 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_train_file_list.json'), 'r') as f: 40 | train_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 41 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_val_file_list.json'), 'r') as f: 42 | val_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 43 | with open(os.path.join(self.root, 'train_test_split', 'shuffled_test_file_list.json'), 'r') as f: 44 | test_ids = set([str(d.split('/')[2]) for d in json.load(f)]) 45 | for item in self.cat: 46 | # print('category', item) 47 | self.meta[item] = [] 48 | dir_point = os.path.join(self.root, self.cat[item]) 49 | fns = sorted(os.listdir(dir_point)) 50 | # print(fns[0][0:-4]) 51 | if split == 'trainval': 52 | fns = [fn for fn in fns if ((fn[0:-4] in train_ids) or (fn[0:-4] in val_ids))] 53 | elif split == 'train': 54 | fns = [fn for fn in fns if fn[0:-4] in train_ids] 55 | elif split == 'val': 56 | fns = [fn for fn in fns if fn[0:-4] in val_ids] 57 | elif split == 'test': 58 | fns = [fn for fn in fns if fn[0:-4] in test_ids] 59 | else: 60 | print('Unknown split: %s. Exiting..' % (split)) 61 | exit(-1) 62 | 63 | # print(fns) 64 | tmp = [] 65 | if split == 'trainval' or split == 'train': 66 | print('+++++++++++++++++++++++++++++++++++++++++++++++++++') 67 | # random.shuffle(fns) 68 | # fns = fns[: int(0.02 * len(fns))] 69 | for idx in range(len(fns)): 70 | if idx%20 == 0: 71 | tmp.append(fns[idx]) 72 | print('+++++++++++++++++++++++++++++++++++++++++++++++++++') 73 | fns = tmp 74 | 75 | # print(os.path.basename(fns)) 76 | for fn in fns: 77 | token = (os.path.splitext(os.path.basename(fn))[0]) 78 | self.meta[item].append(os.path.join(dir_point, token + '.txt')) 79 | 80 | self.datapath = [] 81 | for item in self.cat: 82 | for fn in self.meta[item]: 83 | self.datapath.append((item, fn)) 84 | 85 | self.classes = {} 86 | for i in self.cat.keys(): 87 | self.classes[i] = self.classes_original[i] 88 | 89 | # Mapping from category ('Chair') to a list of int [10,11,12,13] as segmentation labels 90 | self.seg_classes = {'Earphone': [16, 17, 18], 'Motorbike': [30, 31, 32, 33, 34, 35], 'Rocket': [41, 42, 43], 91 | 'Car': [8, 9, 10, 11], 'Laptop': [28, 29], 'Cap': [6, 7], 'Skateboard': [44, 45, 46], 92 | 'Mug': [36, 37], 'Guitar': [19, 20, 21], 'Bag': [4, 5], 'Lamp': [24, 25, 26, 27], 93 | 'Table': [47, 48, 49], 'Airplane': [0, 1, 2, 3], 'Pistol': [38, 39, 40], 94 | 'Chair': [12, 13, 14, 15], 'Knife': [22, 23]} 95 | 96 | # for cat in sorted(self.seg_classes.keys()): 97 | # print(cat, self.seg_classes[cat]) 98 | 99 | self.cache = {} # from index to (point_set, cls, seg) tuple 100 | self.cache_size = 20000 101 | 102 | 103 | def __getitem__(self, index): 104 | if index in self.cache: 105 | ppoint_set, cls, seg = self.cache[index] 106 | else: 107 | fn = self.datapath[index] 108 | cat = self.datapath[index][0] 109 | cls = self.classes[cat] 110 | cls = np.array([cls]).astype(np.int32) 111 | data = np.loadtxt(fn[1]).astype(np.float32) 112 | if not self.normal_channel: 113 | point_set = data[:, 0:3] 114 | else: 115 | point_set = data[:, 0:6] 116 | seg = data[:, -1].astype(np.int32) 117 | if len(self.cache) < self.cache_size: 118 | self.cache[index] = (point_set, cls, seg) 119 | point_set[:, 0:3] = pc_normalize(point_set[:, 0:3]) 120 | 121 | choice = np.random.choice(len(seg), self.npoints, replace=True) 122 | 123 | # resample 124 | point_set = point_set[choice, :] 125 | seg = seg[choice] 126 | 127 | # random shuffle point cloud data 128 | idx = np.arange(len(seg)) 129 | np.random.shuffle(idx) 130 | point_set = point_set[idx, ...] 131 | seg = seg[idx] 132 | 133 | return point_set, cls, seg 134 | 135 | def __len__(self): 136 | return len(self.datapath) 137 | 138 | 139 | 140 | -------------------------------------------------------------------------------- /data_utils/ShapeNetDataLoader.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/ShapeNetDataLoader.pyc -------------------------------------------------------------------------------- /data_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__init__.py -------------------------------------------------------------------------------- /data_utils/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__init__.pyc -------------------------------------------------------------------------------- /data_utils/__pycache__/S3DISDataLoader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__pycache__/S3DISDataLoader.cpython-37.pyc -------------------------------------------------------------------------------- /data_utils/__pycache__/ShapeNetDataLoader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__pycache__/ShapeNetDataLoader.cpython-37.pyc -------------------------------------------------------------------------------- /data_utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /data_utils/__pycache__/indoor3d_util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/data_utils/__pycache__/indoor3d_util.cpython-37.pyc -------------------------------------------------------------------------------- /data_utils/collect_indoor3d_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from indoor3d_util import DATA_PATH, collect_point_label 4 | 5 | BASE_DIR = os.path.dirname(os.path.abspath(__file__)) 6 | ROOT_DIR = os.path.dirname(BASE_DIR) 7 | sys.path.append(BASE_DIR) 8 | 9 | anno_paths = [line.rstrip() for line in open(os.path.join(BASE_DIR, 'meta/anno_paths.txt'))] 10 | anno_paths = [os.path.join(DATA_PATH, p) for p in anno_paths] 11 | 12 | output_folder = os.path.join(ROOT_DIR, 'data/stanford_indoor3d') 13 | if not os.path.exists(output_folder): 14 | os.mkdir(output_folder) 15 | 16 | # Note: there is an extra character in the v1.2 data in Area_5/hallway_6. It's fixed manually. 17 | for anno_path in anno_paths: 18 | print(anno_path) 19 | try: 20 | elements = anno_path.split('/') 21 | out_filename = elements[-3]+'_'+elements[-2]+'.npy' # Area_1_hallway_1.npy 22 | collect_point_label(anno_path, os.path.join(output_folder, out_filename), 'numpy') 23 | except: 24 | print(anno_path, 'ERROR!!') 25 | -------------------------------------------------------------------------------- /data_utils/meta/anno_paths.txt: -------------------------------------------------------------------------------- 1 | Area_1/conferenceRoom_1/Annotations 2 | Area_1/conferenceRoom_2/Annotations 3 | Area_1/copyRoom_1/Annotations 4 | Area_1/hallway_1/Annotations 5 | Area_1/hallway_2/Annotations 6 | Area_1/hallway_3/Annotations 7 | Area_1/hallway_4/Annotations 8 | Area_1/hallway_5/Annotations 9 | Area_1/hallway_6/Annotations 10 | Area_1/hallway_7/Annotations 11 | Area_1/hallway_8/Annotations 12 | Area_1/office_10/Annotations 13 | Area_1/office_11/Annotations 14 | Area_1/office_12/Annotations 15 | Area_1/office_13/Annotations 16 | Area_1/office_14/Annotations 17 | Area_1/office_15/Annotations 18 | Area_1/office_16/Annotations 19 | Area_1/office_17/Annotations 20 | Area_1/office_18/Annotations 21 | Area_1/office_19/Annotations 22 | Area_1/office_1/Annotations 23 | Area_1/office_20/Annotations 24 | Area_1/office_21/Annotations 25 | Area_1/office_22/Annotations 26 | Area_1/office_23/Annotations 27 | Area_1/office_24/Annotations 28 | Area_1/office_25/Annotations 29 | Area_1/office_26/Annotations 30 | Area_1/office_27/Annotations 31 | Area_1/office_28/Annotations 32 | Area_1/office_29/Annotations 33 | Area_1/office_2/Annotations 34 | Area_1/office_30/Annotations 35 | Area_1/office_31/Annotations 36 | Area_1/office_3/Annotations 37 | Area_1/office_4/Annotations 38 | Area_1/office_5/Annotations 39 | Area_1/office_6/Annotations 40 | Area_1/office_7/Annotations 41 | Area_1/office_8/Annotations 42 | Area_1/office_9/Annotations 43 | Area_1/pantry_1/Annotations 44 | Area_1/WC_1/Annotations 45 | Area_2/auditorium_1/Annotations 46 | Area_2/auditorium_2/Annotations 47 | Area_2/conferenceRoom_1/Annotations 48 | Area_2/hallway_10/Annotations 49 | Area_2/hallway_11/Annotations 50 | Area_2/hallway_12/Annotations 51 | Area_2/hallway_1/Annotations 52 | Area_2/hallway_2/Annotations 53 | Area_2/hallway_3/Annotations 54 | Area_2/hallway_4/Annotations 55 | Area_2/hallway_5/Annotations 56 | Area_2/hallway_6/Annotations 57 | Area_2/hallway_7/Annotations 58 | Area_2/hallway_8/Annotations 59 | Area_2/hallway_9/Annotations 60 | Area_2/office_10/Annotations 61 | Area_2/office_11/Annotations 62 | Area_2/office_12/Annotations 63 | Area_2/office_13/Annotations 64 | Area_2/office_14/Annotations 65 | Area_2/office_1/Annotations 66 | Area_2/office_2/Annotations 67 | Area_2/office_3/Annotations 68 | Area_2/office_4/Annotations 69 | Area_2/office_5/Annotations 70 | Area_2/office_6/Annotations 71 | Area_2/office_7/Annotations 72 | Area_2/office_8/Annotations 73 | Area_2/office_9/Annotations 74 | Area_2/storage_1/Annotations 75 | Area_2/storage_2/Annotations 76 | Area_2/storage_3/Annotations 77 | Area_2/storage_4/Annotations 78 | Area_2/storage_5/Annotations 79 | Area_2/storage_6/Annotations 80 | Area_2/storage_7/Annotations 81 | Area_2/storage_8/Annotations 82 | Area_2/storage_9/Annotations 83 | Area_2/WC_1/Annotations 84 | Area_2/WC_2/Annotations 85 | Area_3/conferenceRoom_1/Annotations 86 | Area_3/hallway_1/Annotations 87 | Area_3/hallway_2/Annotations 88 | Area_3/hallway_3/Annotations 89 | Area_3/hallway_4/Annotations 90 | Area_3/hallway_5/Annotations 91 | Area_3/hallway_6/Annotations 92 | Area_3/lounge_1/Annotations 93 | Area_3/lounge_2/Annotations 94 | Area_3/office_10/Annotations 95 | Area_3/office_1/Annotations 96 | Area_3/office_2/Annotations 97 | Area_3/office_3/Annotations 98 | Area_3/office_4/Annotations 99 | Area_3/office_5/Annotations 100 | Area_3/office_6/Annotations 101 | Area_3/office_7/Annotations 102 | Area_3/office_8/Annotations 103 | Area_3/office_9/Annotations 104 | Area_3/storage_1/Annotations 105 | Area_3/storage_2/Annotations 106 | Area_3/WC_1/Annotations 107 | Area_3/WC_2/Annotations 108 | Area_4/conferenceRoom_1/Annotations 109 | Area_4/conferenceRoom_2/Annotations 110 | Area_4/conferenceRoom_3/Annotations 111 | Area_4/hallway_10/Annotations 112 | Area_4/hallway_11/Annotations 113 | Area_4/hallway_12/Annotations 114 | Area_4/hallway_13/Annotations 115 | Area_4/hallway_14/Annotations 116 | Area_4/hallway_1/Annotations 117 | Area_4/hallway_2/Annotations 118 | Area_4/hallway_3/Annotations 119 | Area_4/hallway_4/Annotations 120 | Area_4/hallway_5/Annotations 121 | Area_4/hallway_6/Annotations 122 | Area_4/hallway_7/Annotations 123 | Area_4/hallway_8/Annotations 124 | Area_4/hallway_9/Annotations 125 | Area_4/lobby_1/Annotations 126 | Area_4/lobby_2/Annotations 127 | Area_4/office_10/Annotations 128 | Area_4/office_11/Annotations 129 | Area_4/office_12/Annotations 130 | Area_4/office_13/Annotations 131 | Area_4/office_14/Annotations 132 | Area_4/office_15/Annotations 133 | Area_4/office_16/Annotations 134 | Area_4/office_17/Annotations 135 | Area_4/office_18/Annotations 136 | Area_4/office_19/Annotations 137 | Area_4/office_1/Annotations 138 | Area_4/office_20/Annotations 139 | Area_4/office_21/Annotations 140 | Area_4/office_22/Annotations 141 | Area_4/office_2/Annotations 142 | Area_4/office_3/Annotations 143 | Area_4/office_4/Annotations 144 | Area_4/office_5/Annotations 145 | Area_4/office_6/Annotations 146 | Area_4/office_7/Annotations 147 | Area_4/office_8/Annotations 148 | Area_4/office_9/Annotations 149 | Area_4/storage_1/Annotations 150 | Area_4/storage_2/Annotations 151 | Area_4/storage_3/Annotations 152 | Area_4/storage_4/Annotations 153 | Area_4/WC_1/Annotations 154 | Area_4/WC_2/Annotations 155 | Area_4/WC_3/Annotations 156 | Area_4/WC_4/Annotations 157 | Area_5/conferenceRoom_1/Annotations 158 | Area_5/conferenceRoom_2/Annotations 159 | Area_5/conferenceRoom_3/Annotations 160 | Area_5/hallway_10/Annotations 161 | Area_5/hallway_11/Annotations 162 | Area_5/hallway_12/Annotations 163 | Area_5/hallway_13/Annotations 164 | Area_5/hallway_14/Annotations 165 | Area_5/hallway_15/Annotations 166 | Area_5/hallway_1/Annotations 167 | Area_5/hallway_2/Annotations 168 | Area_5/hallway_3/Annotations 169 | Area_5/hallway_4/Annotations 170 | Area_5/hallway_5/Annotations 171 | Area_5/hallway_6/Annotations 172 | Area_5/hallway_7/Annotations 173 | Area_5/hallway_8/Annotations 174 | Area_5/hallway_9/Annotations 175 | Area_5/lobby_1/Annotations 176 | Area_5/office_10/Annotations 177 | Area_5/office_11/Annotations 178 | Area_5/office_12/Annotations 179 | Area_5/office_13/Annotations 180 | Area_5/office_14/Annotations 181 | Area_5/office_15/Annotations 182 | Area_5/office_16/Annotations 183 | Area_5/office_17/Annotations 184 | Area_5/office_18/Annotations 185 | Area_5/office_19/Annotations 186 | Area_5/office_1/Annotations 187 | Area_5/office_20/Annotations 188 | Area_5/office_21/Annotations 189 | Area_5/office_22/Annotations 190 | Area_5/office_23/Annotations 191 | Area_5/office_24/Annotations 192 | Area_5/office_25/Annotations 193 | Area_5/office_26/Annotations 194 | Area_5/office_27/Annotations 195 | Area_5/office_28/Annotations 196 | Area_5/office_29/Annotations 197 | Area_5/office_2/Annotations 198 | Area_5/office_30/Annotations 199 | Area_5/office_31/Annotations 200 | Area_5/office_32/Annotations 201 | Area_5/office_33/Annotations 202 | Area_5/office_34/Annotations 203 | Area_5/office_35/Annotations 204 | Area_5/office_36/Annotations 205 | Area_5/office_37/Annotations 206 | Area_5/office_38/Annotations 207 | Area_5/office_39/Annotations 208 | Area_5/office_3/Annotations 209 | Area_5/office_40/Annotations 210 | Area_5/office_41/Annotations 211 | Area_5/office_42/Annotations 212 | Area_5/office_4/Annotations 213 | Area_5/office_5/Annotations 214 | Area_5/office_6/Annotations 215 | Area_5/office_7/Annotations 216 | Area_5/office_8/Annotations 217 | Area_5/office_9/Annotations 218 | Area_5/pantry_1/Annotations 219 | Area_5/storage_1/Annotations 220 | Area_5/storage_2/Annotations 221 | Area_5/storage_3/Annotations 222 | Area_5/storage_4/Annotations 223 | Area_5/WC_1/Annotations 224 | Area_5/WC_2/Annotations 225 | Area_6/conferenceRoom_1/Annotations 226 | Area_6/copyRoom_1/Annotations 227 | Area_6/hallway_1/Annotations 228 | Area_6/hallway_2/Annotations 229 | Area_6/hallway_3/Annotations 230 | Area_6/hallway_4/Annotations 231 | Area_6/hallway_5/Annotations 232 | Area_6/hallway_6/Annotations 233 | Area_6/lounge_1/Annotations 234 | Area_6/office_10/Annotations 235 | Area_6/office_11/Annotations 236 | Area_6/office_12/Annotations 237 | Area_6/office_13/Annotations 238 | Area_6/office_14/Annotations 239 | Area_6/office_15/Annotations 240 | Area_6/office_16/Annotations 241 | Area_6/office_17/Annotations 242 | Area_6/office_18/Annotations 243 | Area_6/office_19/Annotations 244 | Area_6/office_1/Annotations 245 | Area_6/office_20/Annotations 246 | Area_6/office_21/Annotations 247 | Area_6/office_22/Annotations 248 | Area_6/office_23/Annotations 249 | Area_6/office_24/Annotations 250 | Area_6/office_25/Annotations 251 | Area_6/office_26/Annotations 252 | Area_6/office_27/Annotations 253 | Area_6/office_28/Annotations 254 | Area_6/office_29/Annotations 255 | Area_6/office_2/Annotations 256 | Area_6/office_30/Annotations 257 | Area_6/office_31/Annotations 258 | Area_6/office_32/Annotations 259 | Area_6/office_33/Annotations 260 | Area_6/office_34/Annotations 261 | Area_6/office_35/Annotations 262 | Area_6/office_36/Annotations 263 | Area_6/office_37/Annotations 264 | Area_6/office_3/Annotations 265 | Area_6/office_4/Annotations 266 | Area_6/office_5/Annotations 267 | Area_6/office_6/Annotations 268 | Area_6/office_7/Annotations 269 | Area_6/office_8/Annotations 270 | Area_6/office_9/Annotations 271 | Area_6/openspace_1/Annotations 272 | Area_6/pantry_1/Annotations 273 | -------------------------------------------------------------------------------- /data_utils/meta/class_names.txt: -------------------------------------------------------------------------------- 1 | ceiling 2 | floor 3 | wall 4 | beam 5 | column 6 | window 7 | door 8 | table 9 | chair 10 | sofa 11 | bookcase 12 | board 13 | clutter 14 | -------------------------------------------------------------------------------- /feature_extractor.py: -------------------------------------------------------------------------------- 1 | # from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 2 | from tools.test_dataloader import TestDataloader 3 | import numpy as np 4 | import os 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from torch.autograd import Variable 9 | import torchvision.models as models 10 | import argparse 11 | import torch.optim as optim 12 | import time 13 | import torchvision.models as models 14 | # from corrnet import SingleViewNet 15 | from models.dgcnn import get_graph_feature 16 | from models.dgcnn import DGCNN 17 | from models.pointnet_part_seg import PointnetPartSeg 18 | # from models.pointnet_part_seg import PointNet_Part 19 | from models.meshnet import MeshNet 20 | from models.SVCNN import Semi3D, SingleViewNet 21 | 22 | def extract(img_net, dgcnn, mesh_net, num_views, split, exp_name): 23 | 24 | dataset = 'ModelNet40' 25 | train_set = TestDataloader(dataset, num_points = 1024, num_views = 2, partition= split) 26 | 27 | print('length of the dataset: ', len(train_set)) 28 | data_loader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=False,num_workers=1) 29 | 30 | start_time = time.time() 31 | 32 | img_feat = np.zeros((len(train_set), 256*4)) 33 | pt_feat = np.zeros((len(train_set), 256*4)) 34 | mesh_feat = np.zeros((len(train_set), 256*4)) 35 | label = np.zeros((len(train_set))) 36 | img_pairs = 18 37 | 38 | 39 | for idx in range(img_pairs): 40 | print('index of image pair: ', idx) 41 | iteration = 0 42 | for data in data_loader: 43 | # pt, img, _, target = data 44 | pt, img, img_v, centers, corners, normals, neighbor_index, target = data 45 | 46 | img = Variable(img).to('cuda') 47 | img_v = Variable(img_v).to('cuda') 48 | 49 | #universal features 50 | img_rfeat, img_zfeat, img_spc = img_net(img) 51 | imgv_rfeat, imgv_zfeat, imgv_spc = img_net(img_v) 52 | 53 | imgF = torch.cat((img_rfeat, img_zfeat, img_spc), dim=1) 54 | imgvF = torch.cat((imgv_rfeat, imgv_zfeat, imgv_spc), dim=1) 55 | 56 | img_feat[iteration,:] = img_feat[iteration,:] + imgF.data.cpu().numpy() + imgvF.data.cpu().numpy() 57 | 58 | # print(iteration) 59 | iteration = iteration + 1 60 | if iteration % 400 == 0: 61 | print('iteration: ', iteration) 62 | 63 | 64 | iteration = 0 65 | for data in data_loader: 66 | # pt, img, _, target = data 67 | pt, img, img_v, centers, corners, normals, neighbor_index, target = data 68 | 69 | target = target[:,0] 70 | 71 | # img = Variable(img).to('cuda') 72 | 73 | pt = Variable(pt).to('cuda') 74 | 75 | centers = Variable(torch.cuda.FloatTensor(centers.cuda())) 76 | corners = Variable(torch.cuda.FloatTensor(corners.cuda())) 77 | normals = Variable(torch.cuda.FloatTensor(normals.cuda())) 78 | neighbor_index = Variable(torch.cuda.LongTensor(neighbor_index.cuda())) 79 | 80 | target = Variable(target).to('cuda') 81 | 82 | pt = pt.permute(0,2,1) 83 | 84 | #This part is for domian spcific features 85 | # cloud_feat, _ = dgcnn(pt) 86 | # feat, _ = img_net(img) 87 | # # feat = img_net(img) 88 | # # feat = feat[:,:,0,0] 89 | # # # # print(feat.size()) 90 | # m_feat, _ = mesh_net(centers, corners, normals, neighbor_index) 91 | 92 | #universal features 93 | cloud_rfeat, cloud_zfeat, cloud_spc = dgcnn(pt) 94 | # img_rfeat, img_zfeat, img_spc = img_net(img) 95 | mesh_rfeat, mesh_zfeat, mesh_specific = mesh_net(centers, corners, normals, neighbor_index) 96 | 97 | # print(cloud_rfeat.size(), cloud_zfeat.size(), cloud_spc.size()) 98 | cloudF = torch.cat((cloud_rfeat, cloud_zfeat, cloud_spc), dim=1) 99 | # imgF = torch.cat((img_rfeat, img_zfeat, img_spc), dim=1) 100 | meshF = torch.cat((mesh_rfeat, mesh_zfeat, mesh_specific), dim=1) 101 | 102 | 103 | # m_feat = m_feat / torch.norm(m_feat) 104 | # cloud_feat = cloud_feat / torch.norm(cloud_feat) 105 | # feat = feat / torch.norm(feat) 106 | 107 | # print(cloud_feat.size(), img_feat.size()) 108 | # img_feat[iteration,:] = imgF.data.cpu().numpy() 109 | pt_feat[iteration,:] = cloudF.data.cpu().numpy() 110 | mesh_feat[iteration,:] = meshF.data.cpu().numpy() 111 | 112 | label[iteration] = target.data.cpu().numpy() 113 | # print(iteration) 114 | iteration = iteration + 1 115 | if iteration % 1000 == 0: 116 | print('iteration: ', iteration) 117 | 118 | #normalize the image features 119 | img_feat = img_feat/(2.0*img_pairs) 120 | 121 | feature_folder = os.path.join('./extracted_features', exp_name) 122 | if not os.path.exists(feature_folder): 123 | os.makedirs(feature_folder) 124 | 125 | img_feat_name = os.path.join('./extracted_features', exp_name, '%s-%s-%s_NtXent_img_feat'%(dataset, split, img_pairs)) 126 | pt_feat_name = os.path.join('./extracted_features', exp_name, '%s-%s_NtXent_cloud1024_feat'%(dataset, split)) 127 | mesh_feat_name = os.path.join('./extracted_features', exp_name, '%s-%s_NtXent_mesh_feat'%(dataset, split)) 128 | label_name = os.path.join('./extracted_features', exp_name, '%s-%s_NtXent_label'%(dataset, split)) 129 | 130 | np.save(img_feat_name, img_feat) 131 | np.save(pt_feat_name, pt_feat) 132 | np.save(mesh_feat_name, mesh_feat) 133 | np.save(label_name, label) 134 | 135 | 136 | def extract_features(args): 137 | 138 | iterations = 10000 139 | num_views = 2 # 1 12 80 140 | # weights_folder = '180view-ModelNet40-pt-mesh-view-pt1024-img112-contrast-specific' 141 | # weights_folder = '180view-ModelNet40-pt-mesh-view-pt1024-img112-contrast-NoSpecific' 142 | weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-Xcontrast-PointMultiAgreement-T095-Fused-Warmup-100percent_xcenter_p10_nt_xw2_aw0_cw9_baseshare_newcenter_l2_100' 143 | 144 | img_net = SingleViewNet(pre_trained = None) 145 | # img_net = torch.nn.DataParallel(img_net) 146 | 147 | img_net_name = './checkpoints/%s/%d-img_net.pkl'%(weights_folder, iterations) 148 | 149 | img_net.load_state_dict(torch.load(img_net_name)['state_dict']) 150 | # img_net = models.resnet18(pretrained=True) 151 | # img_net = list(img_net.children())[:-1] 152 | # img_net = nn.Sequential(*img_net) 153 | # print(img_net) 154 | 155 | dgcnn = DGCNN(args) 156 | dgcnn_name = './checkpoints/%s/%d-pt_net.pkl'%(weights_folder, iterations) 157 | dgcnn.load_state_dict(torch.load(dgcnn_name)['state_dict']) 158 | # pt_net = PointnetPartSeg() 159 | 160 | mesh_net = MeshNet() 161 | mesh_net_name = './checkpoints/%s/%d-mesh_net.pkl'%(weights_folder, iterations) 162 | mesh_net.load_state_dict(torch.load(mesh_net_name)['state_dict']) 163 | 164 | # model = torch.load('./checkpoints/180view-ModelNet40-xentropy-center-pt-img/95000-head_net.pkl') 165 | 166 | img_net = img_net.eval() 167 | dgcnn = dgcnn.eval() 168 | mesh_net = mesh_net.eval() 169 | img_net = img_net.to('cuda') 170 | dgcnn = dgcnn.to('cuda') 171 | mesh_net = mesh_net.to('cuda') 172 | 173 | # print('extracing features for the training split') 174 | extract(img_net, dgcnn, mesh_net, num_views, 'train', exp_name = weights_folder) 175 | 176 | print('extracing features for the testing split') 177 | extract(img_net, dgcnn, mesh_net, num_views, 'test', exp_name = weights_folder) 178 | print('------------------ Al the Features are saved ---------------------------') 179 | 180 | if __name__ == "__main__": 181 | # Training settings 182 | parser = argparse.ArgumentParser(description='RGB and Point Cloud Correspondence') 183 | 184 | parser.add_argument('--no_cuda', type=bool, default=False, 185 | help='enables CUDA training') 186 | parser.add_argument('--seed', type=int, default=1, metavar='S', 187 | help='random seed (default: 1)') 188 | 189 | parser.add_argument('--num_points', type=int, default=1024, 190 | help='num of points to use') 191 | 192 | parser.add_argument('--dropout', type=float, default=0.5, 193 | help='dropout rate') 194 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 195 | help='Dimension of embeddings') 196 | 197 | parser.add_argument('--k', type=int, default=20, metavar='N', 198 | help='Num of nearest neighbors to use') 199 | parser.add_argument('--model_path', type=str, default='', metavar='N', 200 | help='Pretrained model path') 201 | 202 | parser.add_argument('--gpu_id', type=str, default='0', 203 | help='GPU used to train the network') 204 | 205 | args = parser.parse_args() 206 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 207 | extract_features(args) -------------------------------------------------------------------------------- /feature_vis.py: -------------------------------------------------------------------------------- 1 | # from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 2 | from tools.test_dataloader import TestDataloader 3 | import numpy as np 4 | import os 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from torch.autograd import Variable 9 | import torchvision.models as models 10 | import argparse 11 | import torch.optim as optim 12 | import time 13 | import torchvision.models as models 14 | # from corrnet import SingleViewNet 15 | from models.dgcnn import get_graph_feature 16 | from models.dgcnn import DGCNN 17 | from models.pointnet_part_seg import PointnetPartSeg 18 | # from models.pointnet_part_seg import PointNet_Part 19 | from models.meshnet import MeshNet 20 | from models.SVCNN import Semi3D, SingleViewNet, FusionNet 21 | from tools.triplet_dataloader import TripletDataloader 22 | from tqdm import tqdm 23 | from tools.utils import calculate_accuracy 24 | 25 | 26 | 27 | def extract(img_net, dgcnn, mesh_net, fusionnet, num_views, split, exp_name): 28 | 29 | dataset = 'ModelNet40' 30 | 31 | # test_data_set = TripletDataloader(dataset = 'ModelNet40', num_points = 1024, partition='test', perceptange = 10) 32 | test_data_set = TestDataloader(dataset, num_points = 1024, num_views = 2, partition= 'test') 33 | test_data_loader = torch.utils.data.DataLoader(test_data_set, batch_size= 1, shuffle = False, num_workers=1) 34 | 35 | print('length of the dataset: ', len(test_data_set)) 36 | 37 | start_time = time.time() 38 | 39 | img_pred_lst = [] 40 | pt_pred_lst = [] 41 | mesh_pred_lst = [] 42 | fused_pred_lst = [] 43 | mean_pred_lst = [] 44 | geometric_mean_pred_lst = [] 45 | target_lst = [] 46 | 47 | iteration = 0 48 | for data in test_data_loader: 49 | 50 | pt, img, imgV, centers, corners, normals, neighbor_index, target = data #the last one is the target 51 | 52 | pt = Variable(pt).to('cuda') 53 | 54 | img = Variable(img).to('cuda') 55 | imgV = Variable(imgV).to('cuda') 56 | 57 | centers = Variable(torch.cuda.FloatTensor(centers.cuda())) 58 | corners = Variable(torch.cuda.FloatTensor(corners.cuda())) 59 | normals = Variable(torch.cuda.FloatTensor(normals.cuda())) 60 | neighbor_index = Variable(torch.cuda.LongTensor(neighbor_index.cuda())) 61 | 62 | target = torch.squeeze(target, dim = 1) 63 | target = Variable(target).to('cuda') 64 | 65 | pt = pt.permute(0,2,1) 66 | 67 | pt_pred, pt_feat, pt_base = dgcnn(pt) 68 | img_pred, img_feat, img_base = img_net(img, imgV) 69 | mesh_pred, mesh_feat, mesh_base = mesh_net(centers, corners, normals, neighbor_index) 70 | 71 | fused_pred = fusionnet(pt_base, mesh_base, img_base) 72 | 73 | mean_pred = pt_pred + img_pred + mesh_pred 74 | geometric_mean_pred = pt_pred * img_pred * mesh_pred 75 | 76 | 77 | pt_pred_lst.append(pt_pred.data.cpu()) 78 | mesh_pred_lst.append(mesh_pred.data.cpu()) 79 | img_pred_lst.append(img_pred.data.cpu()) 80 | fused_pred_lst.append(fused_pred.data.cpu()) 81 | 82 | mean_pred_lst.append(mean_pred.data.cpu()) 83 | geometric_mean_pred_lst.append(geometric_mean_pred.data.cpu()) 84 | # print(target.size()) 85 | target_lst.append(target.data.cpu()) 86 | 87 | iteration = iteration + 1 88 | if iteration % 400 == 0: 89 | print('iterations: ', iteration) 90 | 91 | pt_pred = torch.cat(pt_pred_lst, dim = 0) 92 | mesh_pred = torch.cat(mesh_pred_lst, dim = 0) 93 | img_pred = torch.cat(img_pred_lst, dim = 0) 94 | fused_pred = torch.cat(fused_pred_lst, dim = 0) 95 | mean_pred = torch.cat(mean_pred_lst, dim = 0) 96 | geometric_mean_pred = torch.cat(geometric_mean_pred_lst, dim = 0) 97 | target = torch.cat(target_lst, dim = 0) 98 | 99 | print('pred size: ', img_pred.size(), mesh_pred.size(), target.size()) 100 | img_acc = calculate_accuracy(img_pred, target) 101 | pt_acc = calculate_accuracy(pt_pred, target) 102 | mesh_acc = calculate_accuracy(mesh_pred, target) 103 | fused_acc = calculate_accuracy(fused_pred, target) 104 | mean_acc = calculate_accuracy(mean_pred, target) 105 | geometric_mean_acc = calculate_accuracy(geometric_mean_pred, target) 106 | 107 | 108 | print('the pt acc: %.4f'%(pt_acc)) 109 | print('the img acc: %.4f'%(img_acc)) 110 | print('the mesh acc: %.4f'%(mesh_acc)) 111 | print('the fused acc: %.4f'%(fused_acc)) 112 | print('the mean acc: %.4f'%(mean_acc)) 113 | print('the geometric mean acc: %.4f'%(geometric_mean_acc)) 114 | 115 | # img_acc = calculate_accuracy(img_pred, target) 116 | # pt_acc = calculate_accuracy(pt_pred, target) 117 | # mesh_acc = calculate_accuracy(mesh_pred, target) 118 | # fused_acc = calculate_accuracy(fused_pred, target) 119 | # # print(img_acc, pt_acc, mesh_acc) 120 | 121 | # img_acc_lst.append(img_acc) 122 | # pt_acc_lst.append(pt_acc) 123 | # mesh_acc_lst.append(mesh_acc) 124 | # fused_acc_lst.append(fused_acc) 125 | 126 | # iteration = iteration + 1 127 | # if iteration % 200 == 0: 128 | # print('the pt acc: ', sum(pt_acc_lst)/len(pt_acc_lst)) 129 | # print('the img acc: ', sum(img_acc_lst)/len(img_acc_lst)) 130 | # print('the mesh acc: ', sum(mesh_acc_lst)/len(mesh_acc_lst)) 131 | # print('the fused acc: ', sum(fused_acc_lst)/len(fused_acc_lst)) 132 | # print('-------------------------------------------------') 133 | 134 | # print('the pt acc: ', sum(pt_acc_lst)/len(pt_acc_lst)) 135 | # print('the img acc: ', sum(img_acc_lst)/len(img_acc_lst)) 136 | # print('the mesh acc: ', sum(mesh_acc_lst)/len(mesh_acc_lst)) 137 | # print('the fused acc: ', sum(fused_acc_lst)/len(fused_acc_lst)) 138 | 139 | 140 | def extract_features(args): 141 | 142 | iterations = 4000 143 | num_views = 2 # 1 12 80 144 | 145 | # weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-Xcontrast-MultiAgreement-T095-Fused-Warmup-2percent' 146 | # weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-2percent-supervised' 147 | weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-Xcontrast-PointMultiAgreement-T095-Fused-Warmup-2percent_xcenter_p10_warmup1_0001' 148 | 149 | img_net = SingleViewNet(pre_trained = True) 150 | # img_net = torch.nn.DataParallel(img_net) 151 | 152 | img_net_name = './checkpoints/%s/%d-img_net.pkl'%(weights_folder, iterations) 153 | img_net.load_state_dict(torch.load(img_net_name)['state_dict']) 154 | 155 | dgcnn = DGCNN(args) 156 | dgcnn_name = './checkpoints/%s/%d-pt_net.pkl'%(weights_folder, iterations) 157 | dgcnn.load_state_dict(torch.load(dgcnn_name)['state_dict']) 158 | 159 | mesh_net = MeshNet() 160 | mesh_net_name = './checkpoints/%s/%d-mesh_net.pkl'%(weights_folder, iterations) 161 | mesh_net.load_state_dict(torch.load(mesh_net_name)['state_dict']) 162 | 163 | fusion_net = FusionNet() 164 | fusion_net_name = './checkpoints/%s/%d-fusion_net.pkl'%(weights_folder, iterations) 165 | fusion_net.load_state_dict(torch.load(fusion_net_name)['state_dict']) 166 | 167 | 168 | img_net = img_net.eval() 169 | dgcnn = dgcnn.eval() 170 | mesh_net = mesh_net.eval() 171 | fusion_net = fusion_net.eval() 172 | 173 | img_net = img_net.to('cuda') 174 | dgcnn = dgcnn.to('cuda') 175 | mesh_net = mesh_net.to('cuda') 176 | fusion_net = fusion_net.to('cuda') 177 | 178 | 179 | print('evaluation for the testing split') 180 | extract(img_net, dgcnn, mesh_net, fusion_net, num_views, 'test', exp_name = weights_folder) 181 | print('------------------ Al the Features are saved ---------------------------') 182 | 183 | if __name__ == "__main__": 184 | # Training settings 185 | parser = argparse.ArgumentParser(description='RGB and Point Cloud Correspondence') 186 | 187 | parser.add_argument('--no_cuda', type=bool, default=False, 188 | help='enables CUDA training') 189 | parser.add_argument('--seed', type=int, default=1, metavar='S', 190 | help='random seed (default: 1)') 191 | 192 | parser.add_argument('--num_points', type=int, default=1024, 193 | help='num of points to use') 194 | 195 | parser.add_argument('--dropout', type=float, default=0.5, 196 | help='dropout rate') 197 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 198 | help='Dimension of embeddings') 199 | 200 | parser.add_argument('--k', type=int, default=20, metavar='N', 201 | help='Num of nearest neighbors to use') 202 | parser.add_argument('--model_path', type=str, default='', metavar='N', 203 | help='Pretrained model path') 204 | 205 | parser.add_argument('--gpu_id', type=str, default='0', 206 | help='GPU used to train the network') 207 | 208 | args = parser.parse_args() 209 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 210 | extract_features(args) -------------------------------------------------------------------------------- /models/.ipynb_checkpoints/SVCNN-checkpoint.py: -------------------------------------------------------------------------------- 1 | # from models.MVCNN import MVCNN 2 | from __future__ import division, absolute_import 3 | from models.dgcnn import DGCNN 4 | from models.resnet import resnet18 5 | from models.meshnet import MeshNet 6 | from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 7 | from tools.utils import calculate_accuracy 8 | import numpy as np 9 | import os 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | import torchvision.models as models 15 | import argparse 16 | import torch.optim as optim 17 | import time 18 | 19 | # The network for extracting image features 20 | class SingleViewNet(nn.Module): 21 | 22 | def __init__(self, pre_trained = None): 23 | super(SingleViewNet, self).__init__() 24 | 25 | print(pre_trained, '**********pre_trained************') 26 | 27 | # if pre_trained: 28 | # self.img_net = torch.load(pre_trained) 29 | # else: 30 | # # print('--------- The ImageNet is loading from ImageNet models --------- ') 31 | 32 | resnet50 = models.resnet50(pretrained = pre_trained) 33 | resnet50 = list(resnet50.children())[:-1] 34 | 35 | resnet18 = models.resnet18(pretrained = pre_trained) 36 | resnet18 = list(resnet18.children())[:-1] 37 | self.img_net = nn.Sequential(*resnet18) 38 | 39 | #head 40 | # self.linear1 = nn.Linear(2048, 512, bias=False) 41 | 42 | self.linear1 = nn.Linear(512, 512, bias=False) 43 | self.bn6 = nn.BatchNorm1d(512) 44 | 45 | # self.bn6 = nn.BatchNorm1d(512) 46 | # self.linear2 = nn.Linear(512, 256) 47 | 48 | # self.pred = nn.Sequential( 49 | # nn.BatchNorm1d(512), 50 | # nn.ReLU(), 51 | # # nn.Dropout(p=0.5), 52 | # nn.Linear(512, 256), 53 | # nn.BatchNorm1d(256), 54 | # nn.ReLU(), 55 | # # nn.Dropout(p=0.5), 56 | # nn.Linear(256, 40) 57 | # ) 58 | 59 | # self.feature = nn.Sequential( 60 | # nn.BatchNorm1d(512), 61 | # nn.ReLU(), 62 | # nn.Linear(512, 256), 63 | # nn.BatchNorm1d(256), 64 | # nn.ReLU(), 65 | # nn.Linear(256, 256) 66 | # ) 67 | # self.groupfeature = nn.Sequential( 68 | # nn.BatchNorm1d(512), 69 | # nn.ReLU(), 70 | # nn.Linear(512, 256), 71 | # # nn.BatchNorm1d(256), 72 | # # nn.ReLU(), 73 | # # nn.Linear(256, 256) 74 | # ) 75 | 76 | 77 | def forward(self, img, imgV): 78 | 79 | img_feat = self.img_net(img) 80 | img_feat = img_feat.squeeze(3) 81 | img_feat = img_feat.squeeze(2) 82 | 83 | img_featV = self.img_net(imgV) 84 | img_featV = img_featV.squeeze(3) 85 | img_featV = img_featV.squeeze(2) 86 | 87 | img_base1 = self.linear1(img_feat) 88 | img_base2 = self.linear1(img_featV) 89 | img_base1 = F.relu(self.bn6(img_base1)) 90 | img_base2 = F.relu(self.bn6(img_base2)) 91 | img_base = torch.max(img_base1, img_base2) 92 | 93 | # img_pred = self.pred(img_base) 94 | # img_feat = self.feature(img_base) 95 | # img_gfeat = self.groupfeature(img_base) 96 | 97 | return img_base 98 | 99 | 100 | # The networks for training contrastive 101 | class FusionHead(nn.Module): 102 | def __init__(self): 103 | super(FusionHead, self).__init__() 104 | 105 | self.pred = nn.Sequential( 106 | # nn.BatchNorm1d(512), 107 | # nn.ReLU(), 108 | # nn.Dropout(p=0.5), 109 | nn.Linear(512, 256), 110 | nn.BatchNorm1d(256), 111 | nn.ReLU(), 112 | # nn.Dropout(p=0.5), 113 | nn.Linear(256, 40) 114 | ) 115 | 116 | self.feature = nn.Sequential( 117 | # nn.BatchNorm1d(512), 118 | # nn.ReLU(), 119 | nn.Linear(512, 256), 120 | nn.BatchNorm1d(256), 121 | nn.ReLU(), 122 | nn.Linear(256, 256) 123 | ) 124 | self.groupfeature = nn.Sequential( 125 | # nn.BatchNorm1d(512), 126 | # nn.ReLU(), 127 | nn.Linear(512, 256), 128 | # nn.BatchNorm1d(256), 129 | # nn.ReLU(), 130 | # nn.Linear(256, 256) 131 | ) 132 | 133 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 134 | def forward(self, point_base, mesh_base, img_base): 135 | 136 | # print(pt_base.size(), mesh_base.size()) 137 | img_pred = self.pred(img_base) 138 | img_feat = self.feature(img_base) 139 | img_gfeat = self.groupfeature(img_base) 140 | 141 | pt_pred = self.pred(point_base) 142 | pt_feat = self.feature(point_base) 143 | pt_gfeat = self.groupfeature(point_base) 144 | 145 | mesh_pred = self.pred(mesh_base) 146 | mesh_feat = self.feature(mesh_base) 147 | mesh_gfeat = self.groupfeature(mesh_base) 148 | 149 | # print(concatenate_feature.size()) 150 | 151 | return img_pred, pt_pred, mesh_pred, img_feat, pt_feat, mesh_feat, img_gfeat, pt_gfeat, mesh_gfeat 152 | 153 | 154 | 155 | class FusionNet(nn.Module): 156 | def __init__(self): 157 | super(FusionNet, self).__init__() 158 | 159 | self.pred = nn.Sequential( 160 | nn.Linear(512 * 3, 256), 161 | nn.BatchNorm1d(256), 162 | nn.ReLU(), 163 | nn.Linear(256, 40) 164 | ) 165 | 166 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 167 | def forward(self, pt_base, mesh_base, img_base): 168 | 169 | # print(pt_base.size(), mesh_base.size()) 170 | concatenate_feature = torch.cat([pt_base, mesh_base, img_base], dim = 1) 171 | 172 | # print(concatenate_feature.size()) 173 | fused_pred = self.pred(concatenate_feature) 174 | 175 | return fused_pred 176 | 177 | 178 | # The networks for training contrastive 179 | class Semi3D(nn.Module): 180 | def __init__(self, img_net, cloud_net, meshnet, fusion_net, fusion_head): 181 | super(Semi3D, self).__init__() 182 | 183 | self.img_net = img_net 184 | self.cloud_net = cloud_net 185 | self.meshnet = meshnet 186 | self.fusion_net = fusion_net 187 | self.fusion_head = fusion_head 188 | 189 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 190 | def forward(self, pt, img, imgV, centers, corners, normals, neighbor_index): 191 | 192 | # get the representations for point cloud data 193 | pt_base = self.cloud_net(pt) # [N,C] 194 | # print(cloud_ris.size(), cloud_zis.size()) 195 | 196 | mesh_base = self.meshnet(centers, corners, normals, neighbor_index) 197 | 198 | # get the representations and the projections 199 | img_base = self.img_net(img, imgV) # [N,C] 200 | 201 | fused_pred = self.fusion_net(pt_base, mesh_base, img_base) 202 | 203 | img_pred, pt_pred, mesh_pred, img_feat, pt_feat, mesh_feat, img_gfeat, pt_gfeat, mesh_gfeat = self.fusion_head(pt_base, mesh_base, img_base) 204 | 205 | 206 | return pt_pred, mesh_pred, img_pred, fused_pred, pt_feat, mesh_feat, img_feat, pt_base, mesh_base, img_base, pt_gfeat, mesh_gfeat, img_gfeat 207 | 208 | 209 | -------------------------------------------------------------------------------- /models/.ipynb_checkpoints/SVCNN_old-checkpoint.py: -------------------------------------------------------------------------------- 1 | # from models.MVCNN import MVCNN 2 | from __future__ import division, absolute_import 3 | from models.dgcnn import DGCNN 4 | from models.resnet import resnet18 5 | from models.meshnet import MeshNet 6 | from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 7 | from tools.utils import calculate_accuracy 8 | import numpy as np 9 | import os 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | import torchvision.models as models 15 | import argparse 16 | import torch.optim as optim 17 | import time 18 | 19 | # The network for extracting image features 20 | class SingleViewNet(nn.Module): 21 | 22 | def __init__(self, pre_trained = None): 23 | super(SingleViewNet, self).__init__() 24 | 25 | # if pre_trained: 26 | # self.img_net = torch.load(pre_trained) 27 | # else: 28 | # # print('--------- The ImageNet is loading from ImageNet models --------- ') 29 | resnet18 = models.resnet18(pretrained = pre_trained) 30 | resnet18 = list(resnet18.children())[:-1] 31 | self.img_net = nn.Sequential(*resnet18) 32 | 33 | #head 34 | self.linear1 = nn.Linear(512, 512, bias=False) 35 | 36 | # self.bn6 = nn.BatchNorm1d(512) 37 | # self.linear2 = nn.Linear(512, 256) 38 | 39 | self.pred = nn.Sequential( 40 | nn.BatchNorm1d(512), 41 | nn.ReLU(), 42 | nn.Linear(512, 256), 43 | nn.BatchNorm1d(256), 44 | nn.ReLU(), 45 | nn.Linear(256, 40) 46 | ) 47 | 48 | self.feature = nn.Sequential( 49 | nn.BatchNorm1d(512), 50 | nn.ReLU(), 51 | nn.Linear(512, 256), 52 | nn.BatchNorm1d(256), 53 | nn.ReLU(), 54 | nn.Linear(256, 256) 55 | ) 56 | 57 | 58 | def forward(self, img, imgV): 59 | 60 | img_feat = self.img_net(img) 61 | img_feat = img_feat.squeeze(3) 62 | img_feat = img_feat.squeeze(2) 63 | 64 | img_featV = self.img_net(imgV) 65 | img_featV = img_featV.squeeze(3) 66 | img_featV = img_featV.squeeze(2) 67 | 68 | img_base1 = self.linear1(img_feat) 69 | img_base2 = self.linear1(img_featV) 70 | 71 | img_base = torch.max(img_base1, img_base2) 72 | 73 | img_pred = self.pred(img_base) 74 | img_feat = self.feature(img_base) 75 | 76 | return img_pred, img_feat, img_base 77 | 78 | # The networks for training contrastive 79 | class Semi3D(nn.Module): 80 | def __init__(self, img_net, cloud_net, meshnet): 81 | super(Semi3D, self).__init__() 82 | 83 | self.img_net = img_net 84 | self.cloud_net = cloud_net 85 | self.meshnet = meshnet 86 | 87 | self.pred = nn.Sequential( 88 | nn.Linear(1024, 256), 89 | nn.BatchNorm1d(256), 90 | nn.ReLU(), 91 | nn.Linear(256, 40) 92 | ) 93 | 94 | 95 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 96 | def forward(self, pt, img, imgV, centers, corners, normals, neighbor_index): 97 | 98 | # get the representations for point cloud data 99 | pt_pred, pt_feat, pt_base = self.cloud_net(pt) # [N,C] 100 | # print(cloud_ris.size(), cloud_zis.size()) 101 | 102 | mesh_pred, mesh_feat, mesh_base = self.meshnet(centers, corners, normals, neighbor_index) 103 | 104 | # get the representations and the projections 105 | img_pred, img_feat, img_base = self.img_net(img, imgV) # [N,C] 106 | 107 | # print(pt_base.size(), mesh_base.size()) 108 | concatenate_feature = torch.cat([pt_base, mesh_base], dim = 1) 109 | 110 | # print(concatenate_feature.size()) 111 | fused_pred = self.pred(concatenate_feature) 112 | 113 | 114 | return pt_pred, mesh_pred, img_pred, fused_pred, pt_feat, mesh_feat, img_feat -------------------------------------------------------------------------------- /models/.ipynb_checkpoints/dgcnn-checkpoint.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: Yue Wang 5 | @Contact: yuewangx@mit.edu 6 | @File: model.py 7 | @Time: 2018/10/13 6:35 PM 8 | """ 9 | 10 | 11 | import os 12 | import sys 13 | import copy 14 | import math 15 | import numpy as np 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | 20 | 21 | def knn(x, k): 22 | inner = -2*torch.matmul(x.transpose(2, 1), x) 23 | xx = torch.sum(x**2, dim=1, keepdim=True) 24 | pairwise_distance = -xx - inner - xx.transpose(2, 1) 25 | 26 | idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k) 27 | return idx 28 | 29 | 30 | def get_graph_feature(x, k=20, idx=None): 31 | batch_size = x.size(0) 32 | num_points = x.size(2) 33 | x = x.view(batch_size, -1, num_points) 34 | if idx is None: 35 | idx = knn(x, k=k) # (batch_size, num_points, k) 36 | device = torch.device('cuda') 37 | 38 | idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points 39 | 40 | idx = idx + idx_base 41 | 42 | idx = idx.view(-1) 43 | 44 | _, num_dims, _ = x.size() 45 | 46 | x = x.transpose(2, 1).contiguous() # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 47 | feature = x.view(batch_size*num_points, -1)[idx, :] 48 | feature = feature.view(batch_size, num_points, k, num_dims) 49 | x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 50 | 51 | feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2) 52 | 53 | return feature 54 | 55 | 56 | class PointNet(nn.Module): 57 | def __init__(self, args, output_channels=40): 58 | super(PointNet, self).__init__() 59 | self.args = args 60 | self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False) 61 | self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 62 | self.conv3 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 63 | self.conv4 = nn.Conv1d(64, 128, kernel_size=1, bias=False) 64 | self.conv5 = nn.Conv1d(128, args.emb_dims, kernel_size=1, bias=False) 65 | self.bn1 = nn.BatchNorm1d(64) 66 | self.bn2 = nn.BatchNorm1d(64) 67 | self.bn3 = nn.BatchNorm1d(64) 68 | self.bn4 = nn.BatchNorm1d(128) 69 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 70 | self.linear1 = nn.Linear(args.emb_dims, 512, bias=False) 71 | self.bn6 = nn.BatchNorm1d(512) 72 | self.dp1 = nn.Dropout() 73 | self.linear2 = nn.Linear(512, output_channels) 74 | 75 | def forward(self, x): 76 | x = F.relu(self.bn1(self.conv1(x))) 77 | x = F.relu(self.bn2(self.conv2(x))) 78 | x = F.relu(self.bn3(self.conv3(x))) 79 | x = F.relu(self.bn4(self.conv4(x))) 80 | x = F.relu(self.bn5(self.conv5(x))) 81 | x = F.adaptive_max_pool1d(x, 1).squeeze() 82 | x = F.relu(self.bn6(self.linear1(x))) 83 | x = self.dp1(x) 84 | x = self.linear2(x) 85 | return x 86 | 87 | 88 | class DGCNN(nn.Module): 89 | def __init__(self, args, output_channels=40): 90 | super(DGCNN, self).__init__() 91 | self.args = args 92 | self.k = args.k 93 | 94 | self.bn1 = nn.BatchNorm2d(64) 95 | self.bn2 = nn.BatchNorm2d(64) 96 | self.bn3 = nn.BatchNorm2d(128) 97 | self.bn4 = nn.BatchNorm2d(256) 98 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 99 | self.bn6 = nn.BatchNorm1d(512) 100 | 101 | self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), 102 | self.bn1, 103 | nn.LeakyReLU(negative_slope=0.2)) 104 | self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), 105 | self.bn2, 106 | nn.LeakyReLU(negative_slope=0.2)) 107 | self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False), 108 | self.bn3, 109 | nn.LeakyReLU(negative_slope=0.2)) 110 | self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False), 111 | self.bn4, 112 | nn.LeakyReLU(negative_slope=0.2)) 113 | self.conv5 = nn.Sequential(nn.Conv1d(512, args.emb_dims, kernel_size=1, bias=False), 114 | self.bn5, 115 | nn.LeakyReLU(negative_slope=0.2)) 116 | 117 | self.linear1 = nn.Linear(args.emb_dims*2, 512, bias=False) 118 | 119 | 120 | # self.dp1 = nn.Dropout(p=args.dropout) 121 | 122 | # self.bn7 = nn.BatchNorm1d(256) 123 | # self.dp2 = nn.Dropout(p=args.dropout) 124 | # self.linear3 = nn.Linear(256, output_channels) 125 | 126 | # self.pred = nn.Sequential( 127 | # nn.BatchNorm1d(512), 128 | # nn.ReLU(), 129 | # # nn.Dropout(p=0.5), 130 | # nn.Linear(512, 256), 131 | # nn.BatchNorm1d(256), 132 | # nn.ReLU(), 133 | # # nn.Dropout(p=0.5), 134 | # nn.Linear(256, 40) 135 | # ) 136 | 137 | # self.feature = nn.Sequential( 138 | # nn.BatchNorm1d(512), 139 | # nn.ReLU(), 140 | # nn.Linear(512, 256), 141 | # nn.BatchNorm1d(256), 142 | # nn.ReLU(), 143 | # nn.Linear(256, 256) 144 | # ) 145 | 146 | # self.groupfeature = nn.Sequential( 147 | # nn.BatchNorm1d(512), 148 | # nn.ReLU(), 149 | # nn.Linear(512, 256), 150 | # # nn.BatchNorm1d(256), 151 | # # nn.ReLU(), 152 | # # nn.Linear(256, 256) 153 | # ) 154 | 155 | 156 | def forward(self, x): 157 | batch_size = x.size(0) 158 | x = get_graph_feature(x, k=self.k) 159 | x = self.conv1(x) 160 | x1 = x.max(dim=-1, keepdim=False)[0] 161 | 162 | x = get_graph_feature(x1, k=self.k) 163 | x = self.conv2(x) 164 | x2 = x.max(dim=-1, keepdim=False)[0] 165 | 166 | x = get_graph_feature(x2, k=self.k) 167 | x = self.conv3(x) 168 | x3 = x.max(dim=-1, keepdim=False)[0] 169 | 170 | x = get_graph_feature(x3, k=self.k) 171 | x = self.conv4(x) 172 | x4 = x.max(dim=-1, keepdim=False)[0] 173 | 174 | concat = torch.cat((x1, x2, x3, x4), dim=1) 175 | 176 | x = self.conv5(concat) 177 | x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) 178 | x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1) 179 | 180 | #This is teh feature after the max pooling and average pooling 181 | cloud_feat = torch.cat((x1, x2), 1) 182 | 183 | #this will be the features for point cloud, make sure the features are not normalized by bn and relu 184 | point_base = self.linear1(cloud_feat) 185 | point_base = F.relu(self.bn6(point_base)) 186 | 187 | # pt_pred = self.pred(point_base) 188 | # pt_feat = self.feature(point_base) 189 | # pt_gfeat = self.groupfeature(point_base) 190 | 191 | return point_base 192 | -------------------------------------------------------------------------------- /models/.ipynb_checkpoints/pointnet-checkpoint.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.parallel 4 | import torch.utils.data 5 | from torch.autograd import Variable 6 | import numpy as np 7 | import torch.nn.functional as F 8 | 9 | 10 | class STN3d(nn.Module): 11 | def __init__(self, channel): 12 | super(STN3d, self).__init__() 13 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 14 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 15 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 16 | self.fc1 = nn.Linear(1024, 512) 17 | self.fc2 = nn.Linear(512, 256) 18 | self.fc3 = nn.Linear(256, 9) 19 | self.relu = nn.ReLU() 20 | 21 | self.bn1 = nn.BatchNorm1d(64) 22 | self.bn2 = nn.BatchNorm1d(128) 23 | self.bn3 = nn.BatchNorm1d(1024) 24 | self.bn4 = nn.BatchNorm1d(512) 25 | self.bn5 = nn.BatchNorm1d(256) 26 | 27 | def forward(self, x): 28 | batchsize = x.size()[0] 29 | x = F.relu(self.bn1(self.conv1(x))) 30 | x = F.relu(self.bn2(self.conv2(x))) 31 | x = F.relu(self.bn3(self.conv3(x))) 32 | x = torch.max(x, 2, keepdim=True)[0] 33 | x = x.view(-1, 1024) 34 | 35 | x = F.relu(self.bn4(self.fc1(x))) 36 | x = F.relu(self.bn5(self.fc2(x))) 37 | x = self.fc3(x) 38 | 39 | iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat( 40 | batchsize, 1) 41 | if x.is_cuda: 42 | iden = iden.cuda() 43 | x = x + iden 44 | x = x.view(-1, 3, 3) 45 | return x 46 | 47 | 48 | class STNkd(nn.Module): 49 | def __init__(self, k=64): 50 | super(STNkd, self).__init__() 51 | self.conv1 = torch.nn.Conv1d(k, 64, 1) 52 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 53 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 54 | self.fc1 = nn.Linear(1024, 512) 55 | self.fc2 = nn.Linear(512, 256) 56 | self.fc3 = nn.Linear(256, k * k) 57 | self.relu = nn.ReLU() 58 | 59 | self.bn1 = nn.BatchNorm1d(64) 60 | self.bn2 = nn.BatchNorm1d(128) 61 | self.bn3 = nn.BatchNorm1d(1024) 62 | self.bn4 = nn.BatchNorm1d(512) 63 | self.bn5 = nn.BatchNorm1d(256) 64 | 65 | self.k = k 66 | 67 | def forward(self, x): 68 | batchsize = x.size()[0] 69 | x = F.relu(self.bn1(self.conv1(x))) 70 | x = F.relu(self.bn2(self.conv2(x))) 71 | x = F.relu(self.bn3(self.conv3(x))) 72 | x = torch.max(x, 2, keepdim=True)[0] 73 | x = x.view(-1, 1024) 74 | 75 | x = F.relu(self.bn4(self.fc1(x))) 76 | x = F.relu(self.bn5(self.fc2(x))) 77 | x = self.fc3(x) 78 | 79 | iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat( 80 | batchsize, 1) 81 | if x.is_cuda: 82 | iden = iden.cuda() 83 | x = x + iden 84 | x = x.view(-1, self.k, self.k) 85 | return x 86 | 87 | 88 | class PointNetEncoder(nn.Module): 89 | def __init__(self, global_feat=True, feature_transform=False, channel=3): 90 | super(PointNetEncoder, self).__init__() 91 | self.stn = STN3d(channel) 92 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 93 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 94 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 95 | self.bn1 = nn.BatchNorm1d(64) 96 | self.bn2 = nn.BatchNorm1d(128) 97 | self.bn3 = nn.BatchNorm1d(1024) 98 | self.global_feat = global_feat 99 | self.feature_transform = feature_transform 100 | if self.feature_transform: 101 | self.fstn = STNkd(k=64) 102 | 103 | def forward(self, x): 104 | B, D, N = x.size() 105 | trans = self.stn(x) 106 | x = x.transpose(2, 1) 107 | if D >3 : 108 | x, feature = x.split(3,dim=2) 109 | x = torch.bmm(x, trans) 110 | if D > 3: 111 | x = torch.cat([x,feature],dim=2) 112 | x = x.transpose(2, 1) 113 | x = F.relu(self.bn1(self.conv1(x))) 114 | 115 | if self.feature_transform: 116 | trans_feat = self.fstn(x) 117 | x = x.transpose(2, 1) 118 | x = torch.bmm(x, trans_feat) 119 | x = x.transpose(2, 1) 120 | else: 121 | trans_feat = None 122 | 123 | pointfeat = x 124 | x = F.relu(self.bn2(self.conv2(x))) 125 | x = self.bn3(self.conv3(x)) 126 | x = torch.max(x, 2, keepdim=True)[0] 127 | x = x.view(-1, 1024) 128 | if self.global_feat: 129 | return x, trans, trans_feat 130 | else: 131 | x = x.view(-1, 1024, 1).repeat(1, 1, N) 132 | return torch.cat([x, pointfeat], 1), trans, trans_feat 133 | 134 | 135 | def feature_transform_reguliarzer(trans): 136 | d = trans.size()[1] 137 | I = torch.eye(d)[None, :, :] 138 | if trans.is_cuda: 139 | I = I.cuda() 140 | loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2, 1) - I), dim=(1, 2))) 141 | return loss 142 | -------------------------------------------------------------------------------- /models/.ipynb_checkpoints/pointnet_fc-checkpoint.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.parallel 5 | import torch.utils.data 6 | import torch.nn.functional as F 7 | from models.pointnet import STN3d, STNkd, feature_transform_reguliarzer 8 | 9 | class PartSeg(nn.Module): 10 | def __init__(self, part_num=50, normal_channel=False): 11 | super(PartSeg, self).__init__() 12 | if normal_channel: 13 | channel = 6 14 | else: 15 | channel = 3 16 | self.part_num = part_num 17 | self.stn = STN3d(channel) 18 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 19 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 20 | self.conv3 = torch.nn.Conv1d(128, 128, 1) 21 | self.conv4 = torch.nn.Conv1d(128, 512, 1) 22 | self.conv5 = torch.nn.Conv1d(512, 2048, 1) 23 | self.bn1 = nn.BatchNorm1d(64) 24 | self.bn2 = nn.BatchNorm1d(128) 25 | self.bn3 = nn.BatchNorm1d(128) 26 | self.bn4 = nn.BatchNorm1d(512) 27 | self.bn5 = nn.BatchNorm1d(2048) 28 | self.fstn = STNkd(k=128) 29 | 30 | def forward(self, point_cloud): 31 | B, D, N = point_cloud.size() 32 | trans = self.stn(point_cloud) 33 | point_cloud = point_cloud.transpose(2, 1) 34 | if D > 3: 35 | point_cloud, feature = point_cloud.split(3, dim=2) 36 | point_cloud = torch.bmm(point_cloud, trans) 37 | if D > 3: 38 | point_cloud = torch.cat([point_cloud, feature], dim=2) 39 | 40 | point_cloud = point_cloud.transpose(2, 1) 41 | 42 | out1 = F.relu(self.bn1(self.conv1(point_cloud))) 43 | out2 = F.relu(self.bn2(self.conv2(out1))) 44 | out3 = F.relu(self.bn3(self.conv3(out2))) 45 | 46 | trans_feat = self.fstn(out3) 47 | x = out3.transpose(2, 1) 48 | net_transformed = torch.bmm(x, trans_feat) 49 | net_transformed = net_transformed.transpose(2, 1) 50 | 51 | out4 = F.relu(self.bn4(self.conv4(net_transformed))) 52 | out5 = self.bn5(self.conv5(out4)) 53 | out_max = torch.max(out5, 2, keepdim=True)[0] 54 | # out_max = out_max.view(-1, 2048) 55 | 56 | # # # out_max = torch.cat([out_max,label.squeeze(1)],1) 57 | expand = out_max.view(-1, 2048, 1).repeat(1, 1, N) 58 | 59 | return torch.cat([expand, out1, out2, out3, out4, out5], 1), out_max 60 | 61 | 62 | class get_model(nn.Module): 63 | 64 | def __init__(self, part_num=50, num_classes=16): 65 | super(get_model, self).__init__() 66 | self.part_num = part_num 67 | self.num_classes = num_classes 68 | self.convs_label = torch.nn.Conv1d(self.num_classes, 64, 1) 69 | self.convs1 = torch.nn.Conv1d(4928+64, 256, 1) 70 | self.convs2 = torch.nn.Conv1d(256, 256, 1) 71 | self.convs3 = torch.nn.Conv1d(256, 128, 1) 72 | self.convs4 = torch.nn.Conv1d(128, self.part_num, 1) 73 | self.bns_label = nn.BatchNorm1d(64) 74 | self.bns1 = nn.BatchNorm1d(256) 75 | self.bns2 = nn.BatchNorm1d(256) 76 | self.bns3 = nn.BatchNorm1d(128) 77 | 78 | def forward(self, pt_feat, pt_cat, label): 79 | B, D, N = pt_cat.size() 80 | # print('size of cat ====', pt_cat.size()) 81 | # print('size of feat ====', pt_feat.size()) 82 | # print('size of label ====', label.size()) 83 | label = label.view(B, self.num_classes, 1) 84 | # print('label ====', label.size()) 85 | label = F.relu(self.bns_label(self.convs_label(label))) 86 | # print('label ====', label.size()) 87 | 88 | # feat_label = torch.cat([pt_feat,label.squeeze()],1) 89 | # print('feat_label ====', feat_label.size()) 90 | # feat_label = feat_label.view(-1, 512+16, 1).repeat(1, 1, N) 91 | 92 | label = label.view(-1, 64, 1).repeat(1, 1, N) 93 | # print('feat_label ====', feat_label.size()) 94 | concat = torch.cat([pt_cat, label], 1) 95 | # print('concat ====', concat.size()) 96 | 97 | net = F.relu(self.bns1(self.convs1(concat))) 98 | net = F.relu(self.bns2(self.convs2(net))) 99 | net = F.relu(self.bns3(self.convs3(net))) 100 | net = self.convs4(net) 101 | net = net.transpose(2, 1).contiguous() 102 | net = F.log_softmax(net.view(-1, self.part_num), dim=-1) 103 | net = net.view(B, N, self.part_num) # [B, N, 50] 104 | 105 | return net 106 | 107 | 108 | class get_loss(torch.nn.Module): 109 | def __init__(self): 110 | super(get_loss, self).__init__() 111 | # self.mat_diff_loss_scale = mat_diff_loss_scale 112 | 113 | def forward(self, pred, target): 114 | loss = F.nll_loss(pred, target) 115 | # mat_diff_loss = feature_transform_reguliarzer(trans_feat) 116 | # total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale 117 | return loss 118 | -------------------------------------------------------------------------------- /models/SVCNN-Copy1.py: -------------------------------------------------------------------------------- 1 | # from models.MVCNN import MVCNN 2 | from __future__ import division, absolute_import 3 | from models.dgcnn import DGCNN 4 | from models.resnet import resnet18 5 | from models.meshnet import MeshNet 6 | from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 7 | from tools.utils import calculate_accuracy 8 | import numpy as np 9 | import os 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | import torchvision.models as models 15 | import argparse 16 | import torch.optim as optim 17 | import time 18 | 19 | # The network for extracting image features 20 | class SingleViewNet(nn.Module): 21 | 22 | def __init__(self, pre_trained = None): 23 | super(SingleViewNet, self).__init__() 24 | 25 | # if pre_trained: 26 | # self.img_net = torch.load(pre_trained) 27 | # else: 28 | # # print('--------- The ImageNet is loading from ImageNet models --------- ') 29 | resnet18 = models.resnet18(pretrained = pre_trained) 30 | resnet18 = list(resnet18.children())[:-1] 31 | self.img_net = nn.Sequential(*resnet18) 32 | 33 | #head 34 | self.linear1 = nn.Linear(512, 512, bias=False) 35 | 36 | # self.bn6 = nn.BatchNorm1d(512) 37 | # self.linear2 = nn.Linear(512, 256) 38 | 39 | self.pred = nn.Sequential( 40 | nn.BatchNorm1d(512), 41 | nn.ReLU(), 42 | # nn.Dropout(p=0.5), 43 | nn.Linear(512, 256), 44 | nn.BatchNorm1d(256), 45 | nn.ReLU(), 46 | # nn.Dropout(p=0.5), 47 | nn.Linear(256, 40) 48 | ) 49 | 50 | self.feature = nn.Sequential( 51 | nn.BatchNorm1d(512), 52 | nn.ReLU(), 53 | nn.Linear(512, 256), 54 | nn.BatchNorm1d(256), 55 | nn.ReLU(), 56 | nn.Linear(256, 256) 57 | ) 58 | self.groupfeature = nn.Sequential( 59 | nn.BatchNorm1d(512), 60 | nn.ReLU(), 61 | nn.Linear(512, 256), 62 | # nn.BatchNorm1d(256), 63 | # nn.ReLU(), 64 | # nn.Linear(256, 256) 65 | ) 66 | 67 | 68 | def forward(self, img, imgV): 69 | 70 | img_feat = self.img_net(img) 71 | img_feat = img_feat.squeeze(3) 72 | img_feat = img_feat.squeeze(2) 73 | 74 | img_featV = self.img_net(imgV) 75 | img_featV = img_featV.squeeze(3) 76 | img_featV = img_featV.squeeze(2) 77 | 78 | img_base1 = self.linear1(img_feat) 79 | img_base2 = self.linear1(img_featV) 80 | 81 | img_base = torch.max(img_base1, img_base2) 82 | 83 | img_pred = self.pred(img_base) 84 | img_feat = self.feature(img_base) 85 | img_gfeat = self.groupfeature(img_base) 86 | 87 | return img_pred, img_feat, img_base, img_gfeat 88 | 89 | 90 | # The networks for training contrastive 91 | class FusionNet(nn.Module): 92 | def __init__(self): 93 | super(FusionNet, self).__init__() 94 | 95 | self.pred = nn.Sequential( 96 | nn.Linear(512 * 3, 256), 97 | nn.BatchNorm1d(256), 98 | nn.ReLU(), 99 | nn.Linear(256, 40) 100 | ) 101 | 102 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 103 | def forward(self, pt_base, mesh_base, img_base): 104 | 105 | # print(pt_base.size(), mesh_base.size()) 106 | concatenate_feature = torch.cat([pt_base, mesh_base, img_base], dim = 1) 107 | 108 | # print(concatenate_feature.size()) 109 | fused_pred = self.pred(concatenate_feature) 110 | 111 | return fused_pred 112 | 113 | 114 | # The networks for training contrastive 115 | class Semi3D(nn.Module): 116 | def __init__(self, img_net, cloud_net, meshnet, fusion_net): 117 | super(Semi3D, self).__init__() 118 | 119 | self.img_net = img_net 120 | self.cloud_net = cloud_net 121 | self.meshnet = meshnet 122 | self.fusion_net = fusion_net 123 | 124 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 125 | def forward(self, pt, img, imgV, centers, corners, normals, neighbor_index): 126 | 127 | # get the representations for point cloud data 128 | pt_pred, pt_feat, pt_base, pt_gfeat = self.cloud_net(pt) # [N,C] 129 | # print(cloud_ris.size(), cloud_zis.size()) 130 | 131 | mesh_pred, mesh_feat, mesh_base, mesh_gfeat = self.meshnet(centers, corners, normals, neighbor_index) 132 | 133 | # get the representations and the projections 134 | img_pred, img_feat, img_base, img_gfeat = self.img_net(img, imgV) # [N,C] 135 | 136 | fused_pred = self.fusion_net(pt_base, mesh_base, img_base) 137 | 138 | 139 | return pt_pred, mesh_pred, img_pred, fused_pred, pt_feat, mesh_feat, img_feat, pt_base, mesh_base, img_base, pt_gfeat, mesh_gfeat, img_gfeat 140 | 141 | 142 | -------------------------------------------------------------------------------- /models/SVCNN.py: -------------------------------------------------------------------------------- 1 | # from models.MVCNN import MVCNN 2 | from __future__ import division, absolute_import 3 | from models.dgcnn import DGCNN 4 | from models.resnet import resnet18 5 | from models.meshnet import MeshNet 6 | from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 7 | from tools.utils import calculate_accuracy 8 | import numpy as np 9 | import os 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | import torchvision.models as models 15 | import argparse 16 | import torch.optim as optim 17 | import time 18 | 19 | # The network for extracting image features 20 | class SingleViewNet(nn.Module): 21 | 22 | def __init__(self, pre_trained = None): 23 | super(SingleViewNet, self).__init__() 24 | 25 | print(pre_trained, '**********pre_trained************') 26 | 27 | # if pre_trained: 28 | # self.img_net = torch.load(pre_trained) 29 | # else: 30 | # # print('--------- The ImageNet is loading from ImageNet models --------- ') 31 | 32 | resnet50 = models.resnet50(pretrained = pre_trained) 33 | resnet50 = list(resnet50.children())[:-1] 34 | 35 | resnet18 = models.resnet18(pretrained = pre_trained) 36 | resnet18 = list(resnet18.children())[:-1] 37 | self.img_net = nn.Sequential(*resnet18) 38 | 39 | #head 40 | # self.linear1 = nn.Linear(2048, 512, bias=False) 41 | 42 | self.linear1 = nn.Linear(512, 512, bias=False) 43 | self.bn6 = nn.BatchNorm1d(512) 44 | 45 | # self.bn6 = nn.BatchNorm1d(512) 46 | # self.linear2 = nn.Linear(512, 256) 47 | 48 | # self.pred = nn.Sequential( 49 | # nn.BatchNorm1d(512), 50 | # nn.ReLU(), 51 | # # nn.Dropout(p=0.5), 52 | # nn.Linear(512, 256), 53 | # nn.BatchNorm1d(256), 54 | # nn.ReLU(), 55 | # # nn.Dropout(p=0.5), 56 | # nn.Linear(256, 40) 57 | # ) 58 | 59 | # self.feature = nn.Sequential( 60 | # nn.BatchNorm1d(512), 61 | # nn.ReLU(), 62 | # nn.Linear(512, 256), 63 | # nn.BatchNorm1d(256), 64 | # nn.ReLU(), 65 | # nn.Linear(256, 256) 66 | # ) 67 | # self.groupfeature = nn.Sequential( 68 | # nn.BatchNorm1d(512), 69 | # nn.ReLU(), 70 | # nn.Linear(512, 256), 71 | # # nn.BatchNorm1d(256), 72 | # # nn.ReLU(), 73 | # # nn.Linear(256, 256) 74 | # ) 75 | 76 | 77 | def forward(self, img, imgV): 78 | 79 | img_feat = self.img_net(img) 80 | img_feat = img_feat.squeeze(3) 81 | img_feat = img_feat.squeeze(2) 82 | 83 | img_featV = self.img_net(imgV) 84 | img_featV = img_featV.squeeze(3) 85 | img_featV = img_featV.squeeze(2) 86 | 87 | img_base1 = self.linear1(img_feat) 88 | img_base2 = self.linear1(img_featV) 89 | img_base1 = F.relu(self.bn6(img_base1)) 90 | img_base2 = F.relu(self.bn6(img_base2)) 91 | img_base = torch.max(img_base1, img_base2) 92 | 93 | # img_pred = self.pred(img_base) 94 | # img_feat = self.feature(img_base) 95 | # img_gfeat = self.groupfeature(img_base) 96 | 97 | return img_base 98 | 99 | 100 | # The networks for training contrastive 101 | class FusionHead(nn.Module): 102 | def __init__(self): 103 | super(FusionHead, self).__init__() 104 | 105 | self.pred = nn.Sequential( 106 | # nn.BatchNorm1d(512), 107 | # nn.ReLU(), 108 | # nn.Dropout(p=0.5), 109 | nn.Linear(512, 256), 110 | nn.BatchNorm1d(256), 111 | nn.ReLU(), 112 | # nn.Dropout(p=0.5), 113 | nn.Linear(256, 40) 114 | ) 115 | 116 | self.feature = nn.Sequential( 117 | # nn.BatchNorm1d(512), 118 | # nn.ReLU(), 119 | nn.Linear(512, 256), 120 | nn.BatchNorm1d(256), 121 | nn.ReLU(), 122 | nn.Linear(256, 256) 123 | ) 124 | self.groupfeature = nn.Sequential( 125 | # nn.BatchNorm1d(512), 126 | # nn.ReLU(), 127 | nn.Linear(512, 256), 128 | # nn.BatchNorm1d(256), 129 | # nn.ReLU(), 130 | # nn.Linear(256, 256) 131 | ) 132 | 133 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 134 | def forward(self, point_base, mesh_base, img_base): 135 | 136 | # print(pt_base.size(), mesh_base.size()) 137 | img_pred = self.pred(img_base) 138 | img_feat = self.feature(img_base) 139 | img_gfeat = self.groupfeature(img_base) 140 | 141 | pt_pred = self.pred(point_base) 142 | pt_feat = self.feature(point_base) 143 | pt_gfeat = self.groupfeature(point_base) 144 | 145 | mesh_pred = self.pred(mesh_base) 146 | mesh_feat = self.feature(mesh_base) 147 | mesh_gfeat = self.groupfeature(mesh_base) 148 | 149 | # print(concatenate_feature.size()) 150 | 151 | return img_pred, pt_pred, mesh_pred, img_feat, pt_feat, mesh_feat, img_gfeat, pt_gfeat, mesh_gfeat 152 | 153 | 154 | 155 | class FusionNet(nn.Module): 156 | def __init__(self): 157 | super(FusionNet, self).__init__() 158 | 159 | self.pred = nn.Sequential( 160 | nn.Linear(512 * 3, 256), 161 | nn.BatchNorm1d(256), 162 | nn.ReLU(), 163 | nn.Linear(256, 40) 164 | ) 165 | 166 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 167 | def forward(self, pt_base, mesh_base, img_base): 168 | 169 | # print(pt_base.size(), mesh_base.size()) 170 | concatenate_feature = torch.cat([pt_base, mesh_base, img_base], dim = 1) 171 | 172 | # print(concatenate_feature.size()) 173 | fused_pred = self.pred(concatenate_feature) 174 | 175 | return fused_pred 176 | 177 | 178 | # The networks for training contrastive 179 | class Semi3D(nn.Module): 180 | def __init__(self, img_net, cloud_net, meshnet, fusion_net, fusion_head): 181 | super(Semi3D, self).__init__() 182 | 183 | self.img_net = img_net 184 | self.cloud_net = cloud_net 185 | self.meshnet = meshnet 186 | self.fusion_net = fusion_net 187 | self.fusion_head = fusion_head 188 | 189 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 190 | def forward(self, pt, img, imgV, centers, corners, normals, neighbor_index): 191 | 192 | # get the representations for point cloud data 193 | pt_base = self.cloud_net(pt) # [N,C] 194 | # print(cloud_ris.size(), cloud_zis.size()) 195 | 196 | mesh_base = self.meshnet(centers, corners, normals, neighbor_index) 197 | 198 | # get the representations and the projections 199 | img_base = self.img_net(img, imgV) # [N,C] 200 | 201 | fused_pred = self.fusion_net(pt_base, mesh_base, img_base) 202 | 203 | img_pred, pt_pred, mesh_pred, img_feat, pt_feat, mesh_feat, img_gfeat, pt_gfeat, mesh_gfeat = self.fusion_head(pt_base, mesh_base, img_base) 204 | 205 | 206 | return pt_pred, mesh_pred, img_pred, fused_pred, pt_feat, mesh_feat, img_feat, pt_base, mesh_base, img_base, pt_gfeat, mesh_gfeat, img_gfeat 207 | 208 | 209 | -------------------------------------------------------------------------------- /models/SVCNN_old.py: -------------------------------------------------------------------------------- 1 | # from models.MVCNN import MVCNN 2 | from __future__ import division, absolute_import 3 | from models.dgcnn import DGCNN 4 | from models.resnet import resnet18 5 | from models.meshnet import MeshNet 6 | from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 7 | from tools.utils import calculate_accuracy 8 | import numpy as np 9 | import os 10 | import torch 11 | import torch.nn as nn 12 | import torch.nn.functional as F 13 | from torch.autograd import Variable 14 | import torchvision.models as models 15 | import argparse 16 | import torch.optim as optim 17 | import time 18 | 19 | # The network for extracting image features 20 | class SingleViewNet(nn.Module): 21 | 22 | def __init__(self, pre_trained = None): 23 | super(SingleViewNet, self).__init__() 24 | 25 | # if pre_trained: 26 | # self.img_net = torch.load(pre_trained) 27 | # else: 28 | # # print('--------- The ImageNet is loading from ImageNet models --------- ') 29 | resnet18 = models.resnet18(pretrained = pre_trained) 30 | resnet18 = list(resnet18.children())[:-1] 31 | self.img_net = nn.Sequential(*resnet18) 32 | 33 | #head 34 | self.linear1 = nn.Linear(512, 512, bias=False) 35 | 36 | # self.bn6 = nn.BatchNorm1d(512) 37 | # self.linear2 = nn.Linear(512, 256) 38 | 39 | self.pred = nn.Sequential( 40 | nn.BatchNorm1d(512), 41 | nn.ReLU(), 42 | nn.Linear(512, 256), 43 | nn.BatchNorm1d(256), 44 | nn.ReLU(), 45 | nn.Linear(256, 40) 46 | ) 47 | 48 | self.feature = nn.Sequential( 49 | nn.BatchNorm1d(512), 50 | nn.ReLU(), 51 | nn.Linear(512, 256), 52 | nn.BatchNorm1d(256), 53 | nn.ReLU(), 54 | nn.Linear(256, 256) 55 | ) 56 | 57 | 58 | def forward(self, img, imgV): 59 | 60 | img_feat = self.img_net(img) 61 | img_feat = img_feat.squeeze(3) 62 | img_feat = img_feat.squeeze(2) 63 | 64 | img_featV = self.img_net(imgV) 65 | img_featV = img_featV.squeeze(3) 66 | img_featV = img_featV.squeeze(2) 67 | 68 | img_base1 = self.linear1(img_feat) 69 | img_base2 = self.linear1(img_featV) 70 | 71 | img_base = torch.max(img_base1, img_base2) 72 | 73 | img_pred = self.pred(img_base) 74 | img_feat = self.feature(img_base) 75 | 76 | return img_pred, img_feat, img_base 77 | 78 | # The networks for training contrastive 79 | class Semi3D(nn.Module): 80 | def __init__(self, img_net, cloud_net, meshnet): 81 | super(Semi3D, self).__init__() 82 | 83 | self.img_net = img_net 84 | self.cloud_net = cloud_net 85 | self.meshnet = meshnet 86 | 87 | self.pred = nn.Sequential( 88 | nn.Linear(1024, 256), 89 | nn.BatchNorm1d(256), 90 | nn.ReLU(), 91 | nn.Linear(256, 40) 92 | ) 93 | 94 | 95 | # pt1, img1, centers1, corners1, normals1, neighbor_index1 96 | def forward(self, pt, img, imgV, centers, corners, normals, neighbor_index): 97 | 98 | # get the representations for point cloud data 99 | pt_pred, pt_feat, pt_base = self.cloud_net(pt) # [N,C] 100 | # print(cloud_ris.size(), cloud_zis.size()) 101 | 102 | mesh_pred, mesh_feat, mesh_base = self.meshnet(centers, corners, normals, neighbor_index) 103 | 104 | # get the representations and the projections 105 | img_pred, img_feat, img_base = self.img_net(img, imgV) # [N,C] 106 | 107 | # print(pt_base.size(), mesh_base.size()) 108 | concatenate_feature = torch.cat([pt_base, mesh_base], dim = 1) 109 | 110 | # print(concatenate_feature.size()) 111 | fused_pred = self.pred(concatenate_feature) 112 | 113 | 114 | return pt_pred, mesh_pred, img_pred, fused_pred, pt_feat, mesh_feat, img_feat -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__init__.py -------------------------------------------------------------------------------- /models/__init__.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__init__.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN_twfview.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN_twfview.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN_twfview.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN_twfview.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/SVCNN_twfview.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/SVCNN_twfview.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/dgcnn.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/dgcnn.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/dgcnn.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/dgcnn.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/dgcnn.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/dgcnn.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/meshnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/meshnet.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/meshnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/meshnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/meshnet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/meshnet.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet_fc.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet_fc.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet_part_seg.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet_part_seg.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet_part_seg.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet_part_seg.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/pointnet_part_seg.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/pointnet_part_seg.cpython-38.pyc -------------------------------------------------------------------------------- /models/__pycache__/resnet.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/resnet.cpython-36.pyc -------------------------------------------------------------------------------- /models/__pycache__/resnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/resnet.cpython-37.pyc -------------------------------------------------------------------------------- /models/__pycache__/resnet.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/__pycache__/resnet.cpython-38.pyc -------------------------------------------------------------------------------- /models/dgcnn-Copy1.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: Yue Wang 5 | @Contact: yuewangx@mit.edu 6 | @File: model.py 7 | @Time: 2018/10/13 6:35 PM 8 | """ 9 | 10 | 11 | import os 12 | import sys 13 | import copy 14 | import math 15 | import numpy as np 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | 20 | 21 | def knn(x, k): 22 | inner = -2*torch.matmul(x.transpose(2, 1), x) 23 | xx = torch.sum(x**2, dim=1, keepdim=True) 24 | pairwise_distance = -xx - inner - xx.transpose(2, 1) 25 | 26 | idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k) 27 | return idx 28 | 29 | 30 | def get_graph_feature(x, k=20, idx=None): 31 | batch_size = x.size(0) 32 | num_points = x.size(2) 33 | x = x.view(batch_size, -1, num_points) 34 | if idx is None: 35 | idx = knn(x, k=k) # (batch_size, num_points, k) 36 | device = torch.device('cuda') 37 | 38 | idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points 39 | 40 | idx = idx + idx_base 41 | 42 | idx = idx.view(-1) 43 | 44 | _, num_dims, _ = x.size() 45 | 46 | x = x.transpose(2, 1).contiguous() # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 47 | feature = x.view(batch_size*num_points, -1)[idx, :] 48 | feature = feature.view(batch_size, num_points, k, num_dims) 49 | x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 50 | 51 | feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2) 52 | 53 | return feature 54 | 55 | 56 | class PointNet(nn.Module): 57 | def __init__(self, args, output_channels=40): 58 | super(PointNet, self).__init__() 59 | self.args = args 60 | self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False) 61 | self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 62 | self.conv3 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 63 | self.conv4 = nn.Conv1d(64, 128, kernel_size=1, bias=False) 64 | self.conv5 = nn.Conv1d(128, args.emb_dims, kernel_size=1, bias=False) 65 | self.bn1 = nn.BatchNorm1d(64) 66 | self.bn2 = nn.BatchNorm1d(64) 67 | self.bn3 = nn.BatchNorm1d(64) 68 | self.bn4 = nn.BatchNorm1d(128) 69 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 70 | self.linear1 = nn.Linear(args.emb_dims, 512, bias=False) 71 | self.bn6 = nn.BatchNorm1d(512) 72 | self.dp1 = nn.Dropout() 73 | self.linear2 = nn.Linear(512, output_channels) 74 | 75 | def forward(self, x): 76 | x = F.relu(self.bn1(self.conv1(x))) 77 | x = F.relu(self.bn2(self.conv2(x))) 78 | x = F.relu(self.bn3(self.conv3(x))) 79 | x = F.relu(self.bn4(self.conv4(x))) 80 | x = F.relu(self.bn5(self.conv5(x))) 81 | x = F.adaptive_max_pool1d(x, 1).squeeze() 82 | x = F.relu(self.bn6(self.linear1(x))) 83 | x = self.dp1(x) 84 | x = self.linear2(x) 85 | return x 86 | 87 | 88 | class DGCNN(nn.Module): 89 | def __init__(self, args, output_channels=40): 90 | super(DGCNN, self).__init__() 91 | self.args = args 92 | self.k = args.k 93 | 94 | self.bn1 = nn.BatchNorm2d(64) 95 | self.bn2 = nn.BatchNorm2d(64) 96 | self.bn3 = nn.BatchNorm2d(128) 97 | self.bn4 = nn.BatchNorm2d(256) 98 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 99 | 100 | self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), 101 | self.bn1, 102 | nn.LeakyReLU(negative_slope=0.2)) 103 | self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), 104 | self.bn2, 105 | nn.LeakyReLU(negative_slope=0.2)) 106 | self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False), 107 | self.bn3, 108 | nn.LeakyReLU(negative_slope=0.2)) 109 | self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False), 110 | self.bn4, 111 | nn.LeakyReLU(negative_slope=0.2)) 112 | self.conv5 = nn.Sequential(nn.Conv1d(512, args.emb_dims, kernel_size=1, bias=False), 113 | self.bn5, 114 | nn.LeakyReLU(negative_slope=0.2)) 115 | 116 | self.linear1 = nn.Linear(args.emb_dims*2, 512, bias=False) 117 | 118 | 119 | # self.dp1 = nn.Dropout(p=args.dropout) 120 | 121 | # self.bn7 = nn.BatchNorm1d(256) 122 | # self.dp2 = nn.Dropout(p=args.dropout) 123 | # self.linear3 = nn.Linear(256, output_channels) 124 | 125 | self.pred = nn.Sequential( 126 | nn.BatchNorm1d(512), 127 | nn.ReLU(), 128 | # nn.Dropout(p=0.5), 129 | nn.Linear(512, 256), 130 | nn.BatchNorm1d(256), 131 | nn.ReLU(), 132 | # nn.Dropout(p=0.5), 133 | nn.Linear(256, 40) 134 | ) 135 | 136 | self.feature = nn.Sequential( 137 | nn.BatchNorm1d(512), 138 | nn.ReLU(), 139 | nn.Linear(512, 256), 140 | nn.BatchNorm1d(256), 141 | nn.ReLU(), 142 | nn.Linear(256, 256) 143 | ) 144 | 145 | self.groupfeature = nn.Sequential( 146 | nn.BatchNorm1d(512), 147 | nn.ReLU(), 148 | nn.Linear(512, 256), 149 | # nn.BatchNorm1d(256), 150 | # nn.ReLU(), 151 | # nn.Linear(256, 256) 152 | ) 153 | 154 | 155 | def forward(self, x): 156 | batch_size = x.size(0) 157 | x = get_graph_feature(x, k=self.k) 158 | x = self.conv1(x) 159 | x1 = x.max(dim=-1, keepdim=False)[0] 160 | 161 | x = get_graph_feature(x1, k=self.k) 162 | x = self.conv2(x) 163 | x2 = x.max(dim=-1, keepdim=False)[0] 164 | 165 | x = get_graph_feature(x2, k=self.k) 166 | x = self.conv3(x) 167 | x3 = x.max(dim=-1, keepdim=False)[0] 168 | 169 | x = get_graph_feature(x3, k=self.k) 170 | x = self.conv4(x) 171 | x4 = x.max(dim=-1, keepdim=False)[0] 172 | 173 | concat = torch.cat((x1, x2, x3, x4), dim=1) 174 | 175 | x = self.conv5(concat) 176 | x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) 177 | x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1) 178 | 179 | #This is teh feature after the max pooling and average pooling 180 | cloud_feat = torch.cat((x1, x2), 1) 181 | 182 | #this will be the features for point cloud, make sure the features are not normalized by bn and relu 183 | point_base = self.linear1(cloud_feat) 184 | 185 | pt_pred = self.pred(point_base) 186 | pt_feat = self.feature(point_base) 187 | pt_gfeat = self.groupfeature(point_base) 188 | 189 | return pt_pred, pt_feat, point_base, pt_gfeat 190 | -------------------------------------------------------------------------------- /models/dgcnn.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ 4 | @Author: Yue Wang 5 | @Contact: yuewangx@mit.edu 6 | @File: model.py 7 | @Time: 2018/10/13 6:35 PM 8 | """ 9 | 10 | 11 | import os 12 | import sys 13 | import copy 14 | import math 15 | import numpy as np 16 | import torch 17 | import torch.nn as nn 18 | import torch.nn.functional as F 19 | 20 | 21 | def knn(x, k): 22 | inner = -2*torch.matmul(x.transpose(2, 1), x) 23 | xx = torch.sum(x**2, dim=1, keepdim=True) 24 | pairwise_distance = -xx - inner - xx.transpose(2, 1) 25 | 26 | idx = pairwise_distance.topk(k=k, dim=-1)[1] # (batch_size, num_points, k) 27 | return idx 28 | 29 | 30 | def get_graph_feature(x, k=20, idx=None): 31 | batch_size = x.size(0) 32 | num_points = x.size(2) 33 | x = x.view(batch_size, -1, num_points) 34 | if idx is None: 35 | idx = knn(x, k=k) # (batch_size, num_points, k) 36 | device = torch.device('cuda') 37 | 38 | idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1)*num_points 39 | 40 | idx = idx + idx_base 41 | 42 | idx = idx.view(-1) 43 | 44 | _, num_dims, _ = x.size() 45 | 46 | x = x.transpose(2, 1).contiguous() # (batch_size, num_points, num_dims) -> (batch_size*num_points, num_dims) # batch_size * num_points * k + range(0, batch_size*num_points) 47 | feature = x.view(batch_size*num_points, -1)[idx, :] 48 | feature = feature.view(batch_size, num_points, k, num_dims) 49 | x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1) 50 | 51 | feature = torch.cat((feature-x, x), dim=3).permute(0, 3, 1, 2) 52 | 53 | return feature 54 | 55 | 56 | class PointNet(nn.Module): 57 | def __init__(self, args, output_channels=40): 58 | super(PointNet, self).__init__() 59 | self.args = args 60 | self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False) 61 | self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 62 | self.conv3 = nn.Conv1d(64, 64, kernel_size=1, bias=False) 63 | self.conv4 = nn.Conv1d(64, 128, kernel_size=1, bias=False) 64 | self.conv5 = nn.Conv1d(128, args.emb_dims, kernel_size=1, bias=False) 65 | self.bn1 = nn.BatchNorm1d(64) 66 | self.bn2 = nn.BatchNorm1d(64) 67 | self.bn3 = nn.BatchNorm1d(64) 68 | self.bn4 = nn.BatchNorm1d(128) 69 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 70 | self.linear1 = nn.Linear(args.emb_dims, 512, bias=False) 71 | self.bn6 = nn.BatchNorm1d(512) 72 | self.dp1 = nn.Dropout() 73 | self.linear2 = nn.Linear(512, output_channels) 74 | 75 | def forward(self, x): 76 | x = F.relu(self.bn1(self.conv1(x))) 77 | x = F.relu(self.bn2(self.conv2(x))) 78 | x = F.relu(self.bn3(self.conv3(x))) 79 | x = F.relu(self.bn4(self.conv4(x))) 80 | x = F.relu(self.bn5(self.conv5(x))) 81 | x = F.adaptive_max_pool1d(x, 1).squeeze() 82 | x = F.relu(self.bn6(self.linear1(x))) 83 | x = self.dp1(x) 84 | x = self.linear2(x) 85 | return x 86 | 87 | 88 | class DGCNN(nn.Module): 89 | def __init__(self, args, output_channels=40): 90 | super(DGCNN, self).__init__() 91 | self.args = args 92 | self.k = args.k 93 | 94 | self.bn1 = nn.BatchNorm2d(64) 95 | self.bn2 = nn.BatchNorm2d(64) 96 | self.bn3 = nn.BatchNorm2d(128) 97 | self.bn4 = nn.BatchNorm2d(256) 98 | self.bn5 = nn.BatchNorm1d(args.emb_dims) 99 | self.bn6 = nn.BatchNorm1d(512) 100 | 101 | self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False), 102 | self.bn1, 103 | nn.LeakyReLU(negative_slope=0.2)) 104 | self.conv2 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False), 105 | self.bn2, 106 | nn.LeakyReLU(negative_slope=0.2)) 107 | self.conv3 = nn.Sequential(nn.Conv2d(64*2, 128, kernel_size=1, bias=False), 108 | self.bn3, 109 | nn.LeakyReLU(negative_slope=0.2)) 110 | self.conv4 = nn.Sequential(nn.Conv2d(128*2, 256, kernel_size=1, bias=False), 111 | self.bn4, 112 | nn.LeakyReLU(negative_slope=0.2)) 113 | self.conv5 = nn.Sequential(nn.Conv1d(512, args.emb_dims, kernel_size=1, bias=False), 114 | self.bn5, 115 | nn.LeakyReLU(negative_slope=0.2)) 116 | 117 | self.linear1 = nn.Linear(args.emb_dims*2, 512, bias=False) 118 | 119 | 120 | # self.dp1 = nn.Dropout(p=args.dropout) 121 | 122 | # self.bn7 = nn.BatchNorm1d(256) 123 | # self.dp2 = nn.Dropout(p=args.dropout) 124 | # self.linear3 = nn.Linear(256, output_channels) 125 | 126 | # self.pred = nn.Sequential( 127 | # nn.BatchNorm1d(512), 128 | # nn.ReLU(), 129 | # # nn.Dropout(p=0.5), 130 | # nn.Linear(512, 256), 131 | # nn.BatchNorm1d(256), 132 | # nn.ReLU(), 133 | # # nn.Dropout(p=0.5), 134 | # nn.Linear(256, 40) 135 | # ) 136 | 137 | # self.feature = nn.Sequential( 138 | # nn.BatchNorm1d(512), 139 | # nn.ReLU(), 140 | # nn.Linear(512, 256), 141 | # nn.BatchNorm1d(256), 142 | # nn.ReLU(), 143 | # nn.Linear(256, 256) 144 | # ) 145 | 146 | # self.groupfeature = nn.Sequential( 147 | # nn.BatchNorm1d(512), 148 | # nn.ReLU(), 149 | # nn.Linear(512, 256), 150 | # # nn.BatchNorm1d(256), 151 | # # nn.ReLU(), 152 | # # nn.Linear(256, 256) 153 | # ) 154 | 155 | 156 | def forward(self, x): 157 | batch_size = x.size(0) 158 | x = get_graph_feature(x, k=self.k) 159 | x = self.conv1(x) 160 | x1 = x.max(dim=-1, keepdim=False)[0] 161 | 162 | x = get_graph_feature(x1, k=self.k) 163 | x = self.conv2(x) 164 | x2 = x.max(dim=-1, keepdim=False)[0] 165 | 166 | x = get_graph_feature(x2, k=self.k) 167 | x = self.conv3(x) 168 | x3 = x.max(dim=-1, keepdim=False)[0] 169 | 170 | x = get_graph_feature(x3, k=self.k) 171 | x = self.conv4(x) 172 | x4 = x.max(dim=-1, keepdim=False)[0] 173 | 174 | concat = torch.cat((x1, x2, x3, x4), dim=1) 175 | 176 | x = self.conv5(concat) 177 | x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1) 178 | x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1) 179 | 180 | #This is teh feature after the max pooling and average pooling 181 | cloud_feat = torch.cat((x1, x2), 1) 182 | 183 | #this will be the features for point cloud, make sure the features are not normalized by bn and relu 184 | point_base = self.linear1(cloud_feat) 185 | point_base = F.relu(self.bn6(point_base)) 186 | 187 | # pt_pred = self.pred(point_base) 188 | # pt_feat = self.feature(point_base) 189 | # pt_gfeat = self.groupfeature(point_base) 190 | 191 | return point_base 192 | -------------------------------------------------------------------------------- /models/dgcnn.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/models/dgcnn.pyc -------------------------------------------------------------------------------- /models/pointnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.parallel 4 | import torch.utils.data 5 | from torch.autograd import Variable 6 | import numpy as np 7 | import torch.nn.functional as F 8 | 9 | 10 | class STN3d(nn.Module): 11 | def __init__(self, channel): 12 | super(STN3d, self).__init__() 13 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 14 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 15 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 16 | self.fc1 = nn.Linear(1024, 512) 17 | self.fc2 = nn.Linear(512, 256) 18 | self.fc3 = nn.Linear(256, 9) 19 | self.relu = nn.ReLU() 20 | 21 | self.bn1 = nn.BatchNorm1d(64) 22 | self.bn2 = nn.BatchNorm1d(128) 23 | self.bn3 = nn.BatchNorm1d(1024) 24 | self.bn4 = nn.BatchNorm1d(512) 25 | self.bn5 = nn.BatchNorm1d(256) 26 | 27 | def forward(self, x): 28 | batchsize = x.size()[0] 29 | x = F.relu(self.bn1(self.conv1(x))) 30 | x = F.relu(self.bn2(self.conv2(x))) 31 | x = F.relu(self.bn3(self.conv3(x))) 32 | x = torch.max(x, 2, keepdim=True)[0] 33 | x = x.view(-1, 1024) 34 | 35 | x = F.relu(self.bn4(self.fc1(x))) 36 | x = F.relu(self.bn5(self.fc2(x))) 37 | x = self.fc3(x) 38 | 39 | iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat( 40 | batchsize, 1) 41 | if x.is_cuda: 42 | iden = iden.cuda() 43 | x = x + iden 44 | x = x.view(-1, 3, 3) 45 | return x 46 | 47 | 48 | class STNkd(nn.Module): 49 | def __init__(self, k=64): 50 | super(STNkd, self).__init__() 51 | self.conv1 = torch.nn.Conv1d(k, 64, 1) 52 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 53 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 54 | self.fc1 = nn.Linear(1024, 512) 55 | self.fc2 = nn.Linear(512, 256) 56 | self.fc3 = nn.Linear(256, k * k) 57 | self.relu = nn.ReLU() 58 | 59 | self.bn1 = nn.BatchNorm1d(64) 60 | self.bn2 = nn.BatchNorm1d(128) 61 | self.bn3 = nn.BatchNorm1d(1024) 62 | self.bn4 = nn.BatchNorm1d(512) 63 | self.bn5 = nn.BatchNorm1d(256) 64 | 65 | self.k = k 66 | 67 | def forward(self, x): 68 | batchsize = x.size()[0] 69 | x = F.relu(self.bn1(self.conv1(x))) 70 | x = F.relu(self.bn2(self.conv2(x))) 71 | x = F.relu(self.bn3(self.conv3(x))) 72 | x = torch.max(x, 2, keepdim=True)[0] 73 | x = x.view(-1, 1024) 74 | 75 | x = F.relu(self.bn4(self.fc1(x))) 76 | x = F.relu(self.bn5(self.fc2(x))) 77 | x = self.fc3(x) 78 | 79 | iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat( 80 | batchsize, 1) 81 | if x.is_cuda: 82 | iden = iden.cuda() 83 | x = x + iden 84 | x = x.view(-1, self.k, self.k) 85 | return x 86 | 87 | 88 | class PointNetEncoder(nn.Module): 89 | def __init__(self, global_feat=True, feature_transform=False, channel=3): 90 | super(PointNetEncoder, self).__init__() 91 | self.stn = STN3d(channel) 92 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 93 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 94 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 95 | self.bn1 = nn.BatchNorm1d(64) 96 | self.bn2 = nn.BatchNorm1d(128) 97 | self.bn3 = nn.BatchNorm1d(1024) 98 | self.global_feat = global_feat 99 | self.feature_transform = feature_transform 100 | if self.feature_transform: 101 | self.fstn = STNkd(k=64) 102 | 103 | def forward(self, x): 104 | B, D, N = x.size() 105 | trans = self.stn(x) 106 | x = x.transpose(2, 1) 107 | if D >3 : 108 | x, feature = x.split(3,dim=2) 109 | x = torch.bmm(x, trans) 110 | if D > 3: 111 | x = torch.cat([x,feature],dim=2) 112 | x = x.transpose(2, 1) 113 | x = F.relu(self.bn1(self.conv1(x))) 114 | 115 | if self.feature_transform: 116 | trans_feat = self.fstn(x) 117 | x = x.transpose(2, 1) 118 | x = torch.bmm(x, trans_feat) 119 | x = x.transpose(2, 1) 120 | else: 121 | trans_feat = None 122 | 123 | pointfeat = x 124 | x = F.relu(self.bn2(self.conv2(x))) 125 | x = self.bn3(self.conv3(x)) 126 | x = torch.max(x, 2, keepdim=True)[0] 127 | x = x.view(-1, 1024) 128 | if self.global_feat: 129 | return x, trans, trans_feat 130 | else: 131 | x = x.view(-1, 1024, 1).repeat(1, 1, N) 132 | return torch.cat([x, pointfeat], 1), trans, trans_feat 133 | 134 | 135 | def feature_transform_reguliarzer(trans): 136 | d = trans.size()[1] 137 | I = torch.eye(d)[None, :, :] 138 | if trans.is_cuda: 139 | I = I.cuda() 140 | loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2, 1) - I), dim=(1, 2))) 141 | return loss 142 | -------------------------------------------------------------------------------- /models/pointnet_fc.py: -------------------------------------------------------------------------------- 1 | 2 | import torch 3 | import torch.nn as nn 4 | import torch.nn.parallel 5 | import torch.utils.data 6 | import torch.nn.functional as F 7 | from models.pointnet import STN3d, STNkd, feature_transform_reguliarzer 8 | 9 | class PartSeg(nn.Module): 10 | def __init__(self, part_num=50, normal_channel=False): 11 | super(PartSeg, self).__init__() 12 | if normal_channel: 13 | channel = 6 14 | else: 15 | channel = 3 16 | self.part_num = part_num 17 | self.stn = STN3d(channel) 18 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 19 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 20 | self.conv3 = torch.nn.Conv1d(128, 128, 1) 21 | self.conv4 = torch.nn.Conv1d(128, 512, 1) 22 | self.conv5 = torch.nn.Conv1d(512, 2048, 1) 23 | self.bn1 = nn.BatchNorm1d(64) 24 | self.bn2 = nn.BatchNorm1d(128) 25 | self.bn3 = nn.BatchNorm1d(128) 26 | self.bn4 = nn.BatchNorm1d(512) 27 | self.bn5 = nn.BatchNorm1d(2048) 28 | self.fstn = STNkd(k=128) 29 | 30 | def forward(self, point_cloud): 31 | B, D, N = point_cloud.size() 32 | trans = self.stn(point_cloud) 33 | point_cloud = point_cloud.transpose(2, 1) 34 | if D > 3: 35 | point_cloud, feature = point_cloud.split(3, dim=2) 36 | point_cloud = torch.bmm(point_cloud, trans) 37 | if D > 3: 38 | point_cloud = torch.cat([point_cloud, feature], dim=2) 39 | 40 | point_cloud = point_cloud.transpose(2, 1) 41 | 42 | out1 = F.relu(self.bn1(self.conv1(point_cloud))) 43 | out2 = F.relu(self.bn2(self.conv2(out1))) 44 | out3 = F.relu(self.bn3(self.conv3(out2))) 45 | 46 | trans_feat = self.fstn(out3) 47 | x = out3.transpose(2, 1) 48 | net_transformed = torch.bmm(x, trans_feat) 49 | net_transformed = net_transformed.transpose(2, 1) 50 | 51 | out4 = F.relu(self.bn4(self.conv4(net_transformed))) 52 | out5 = self.bn5(self.conv5(out4)) 53 | out_max = torch.max(out5, 2, keepdim=True)[0] 54 | # out_max = out_max.view(-1, 2048) 55 | 56 | # # # out_max = torch.cat([out_max,label.squeeze(1)],1) 57 | expand = out_max.view(-1, 2048, 1).repeat(1, 1, N) 58 | 59 | return torch.cat([expand, out1, out2, out3, out4, out5], 1), out_max 60 | 61 | 62 | class get_model(nn.Module): 63 | 64 | def __init__(self, part_num=50, num_classes=16): 65 | super(get_model, self).__init__() 66 | self.part_num = part_num 67 | self.num_classes = num_classes 68 | self.convs_label = torch.nn.Conv1d(self.num_classes, 64, 1) 69 | self.convs1 = torch.nn.Conv1d(4928+64, 256, 1) 70 | self.convs2 = torch.nn.Conv1d(256, 256, 1) 71 | self.convs3 = torch.nn.Conv1d(256, 128, 1) 72 | self.convs4 = torch.nn.Conv1d(128, self.part_num, 1) 73 | self.bns_label = nn.BatchNorm1d(64) 74 | self.bns1 = nn.BatchNorm1d(256) 75 | self.bns2 = nn.BatchNorm1d(256) 76 | self.bns3 = nn.BatchNorm1d(128) 77 | 78 | def forward(self, pt_feat, pt_cat, label): 79 | B, D, N = pt_cat.size() 80 | # print('size of cat ====', pt_cat.size()) 81 | # print('size of feat ====', pt_feat.size()) 82 | # print('size of label ====', label.size()) 83 | label = label.view(B, self.num_classes, 1) 84 | # print('label ====', label.size()) 85 | label = F.relu(self.bns_label(self.convs_label(label))) 86 | # print('label ====', label.size()) 87 | 88 | # feat_label = torch.cat([pt_feat,label.squeeze()],1) 89 | # print('feat_label ====', feat_label.size()) 90 | # feat_label = feat_label.view(-1, 512+16, 1).repeat(1, 1, N) 91 | 92 | label = label.view(-1, 64, 1).repeat(1, 1, N) 93 | # print('feat_label ====', feat_label.size()) 94 | concat = torch.cat([pt_cat, label], 1) 95 | # print('concat ====', concat.size()) 96 | 97 | net = F.relu(self.bns1(self.convs1(concat))) 98 | net = F.relu(self.bns2(self.convs2(net))) 99 | net = F.relu(self.bns3(self.convs3(net))) 100 | net = self.convs4(net) 101 | net = net.transpose(2, 1).contiguous() 102 | net = F.log_softmax(net.view(-1, self.part_num), dim=-1) 103 | net = net.view(B, N, self.part_num) # [B, N, 50] 104 | 105 | return net 106 | 107 | 108 | class get_loss(torch.nn.Module): 109 | def __init__(self): 110 | super(get_loss, self).__init__() 111 | # self.mat_diff_loss_scale = mat_diff_loss_scale 112 | 113 | def forward(self, pred, target): 114 | loss = F.nll_loss(pred, target) 115 | # mat_diff_loss = feature_transform_reguliarzer(trans_feat) 116 | # total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale 117 | return loss 118 | -------------------------------------------------------------------------------- /models/pointnet_part_seg.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.parallel 4 | import torch.utils.data 5 | import torch.nn.functional as F 6 | from models.pointnet import STN3d, STNkd, feature_transform_reguliarzer 7 | 8 | 9 | class PointnetPartSeg(nn.Module): 10 | def __init__(self, part_num=50, normal_channel=False): 11 | super(PointnetPartSeg, self).__init__() 12 | if normal_channel: 13 | channel = 6 14 | else: 15 | channel = 3 16 | self.part_num = part_num 17 | self.stn = STN3d(channel) 18 | self.conv1 = torch.nn.Conv1d(channel, 64, 1) 19 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 20 | self.conv3 = torch.nn.Conv1d(128, 128, 1) 21 | self.conv4 = torch.nn.Conv1d(128, 512, 1) 22 | self.conv5 = torch.nn.Conv1d(512, 2048, 1) 23 | self.bn1 = nn.BatchNorm1d(64) 24 | self.bn2 = nn.BatchNorm1d(128) 25 | self.bn3 = nn.BatchNorm1d(128) 26 | self.bn4 = nn.BatchNorm1d(512) 27 | self.bn5 = nn.BatchNorm1d(2048) 28 | self.fstn = STNkd(k=128) 29 | 30 | # self.linear1 = torch.nn.Linear(2048, 512) 31 | # self.bn6 = nn.BatchNorm1d(512) 32 | 33 | # self.linear2 = torch.nn.Linear(512, 256) 34 | 35 | # self.convs3 = torch.nn.Conv1d(256, 128, 1) 36 | # self.convs4 = torch.nn.Conv1d(128, part_num, 1) 37 | 38 | # self.bns2 = nn.BatchNorm1d(256) 39 | # self.bns3 = nn.BatchNorm1d(128) 40 | 41 | def forward(self, point_cloud): 42 | B, D, N = point_cloud.size() 43 | trans = self.stn(point_cloud) 44 | point_cloud = point_cloud.transpose(2, 1) 45 | if D > 3: 46 | point_cloud, feature = point_cloud.split(3, dim=2) 47 | point_cloud = torch.bmm(point_cloud, trans) 48 | if D > 3: 49 | point_cloud = torch.cat([point_cloud, feature], dim=2) 50 | 51 | point_cloud = point_cloud.transpose(2, 1) 52 | 53 | out1 = F.relu(self.bn1(self.conv1(point_cloud))) 54 | out2 = F.relu(self.bn2(self.conv2(out1))) 55 | out3 = F.relu(self.bn3(self.conv3(out2))) 56 | 57 | trans_feat = self.fstn(out3) 58 | x = out3.transpose(2, 1) 59 | net_transformed = torch.bmm(x, trans_feat) 60 | net_transformed = net_transformed.transpose(2, 1) 61 | 62 | out4 = F.relu(self.bn4(self.conv4(net_transformed))) 63 | out5 = self.bn5(self.conv5(out4)) 64 | out_max = torch.max(out5, 2, keepdim=True)[0] 65 | # out_max = out_max.view(-1, 2048) 66 | 67 | # # # out_max = torch.cat([out_max,label.squeeze(1)],1) 68 | expand = out_max.view(-1, 2048, 1).repeat(1, 1, N) 69 | # # cloud_feat = torch.cat([expand, out1, out2, out3, out4, out5], 1) 70 | 71 | # cloud_feat = self.linear1(out_max) 72 | 73 | # # print(cloud_feat.size()) 74 | # #head 75 | # cloud_out = self.linear2(F.relu(self.bn6(cloud_feat))) 76 | # # print(cloud_out.size()) 77 | # return cloud_feat, cloud_out 78 | return torch.cat([expand, out1, out2, out3, out4, out5], 1), out_max 79 | 80 | 81 | # net = F.relu(self.bns1(self.convs1(concat))) 82 | # net = F.relu(self.bns2(self.convs2(net))) 83 | # net = F.relu(self.bns3(self.convs3(net))) 84 | # net = self.convs4(net) 85 | # net = net.transpose(2, 1).contiguous() 86 | # net = F.log_softmax(net.view(-1, self.part_num), dim=-1) 87 | # net = net.view(B, N, self.part_num) # [B, N, 50] 88 | 89 | # return net, trans_feat 90 | 91 | 92 | 93 | # class get_loss(torch.nn.Module): 94 | # def __init__(self, mat_diff_loss_scale=0.001): 95 | # super(get_loss, self).__init__() 96 | # self.mat_diff_loss_scale = mat_diff_loss_scale 97 | 98 | # def forward(self, pred, target, trans_feat): 99 | # loss = F.nll_loss(pred, target) 100 | # mat_diff_loss = feature_transform_reguliarzer(trans_feat) 101 | # total_loss = loss + mat_diff_loss * self.mat_diff_loss_scale 102 | # return total_loss -------------------------------------------------------------------------------- /nt_xent.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | class NTXentLoss(torch.nn.Module): 5 | 6 | def __init__(self, device, batch_size, temperature, use_cosine_similarity): 7 | super(NTXentLoss, self).__init__() 8 | self.batch_size = batch_size 9 | self.temperature = temperature 10 | self.device = device 11 | self.softmax = torch.nn.Softmax(dim=-1) 12 | self.mask_samples_from_same_repr = self._get_correlated_mask().type(torch.bool) 13 | self.similarity_function = self._get_similarity_function(use_cosine_similarity) 14 | self.criterion = torch.nn.CrossEntropyLoss(reduction="sum") 15 | 16 | def _get_similarity_function(self, use_cosine_similarity): 17 | if use_cosine_similarity: 18 | self._cosine_similarity = torch.nn.CosineSimilarity(dim=-1) 19 | return self._cosine_simililarity 20 | else: 21 | return self._dot_simililarity 22 | 23 | def _get_correlated_mask(self): 24 | diag = np.eye(2 * self.batch_size) 25 | l1 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=-self.batch_size) 26 | l2 = np.eye((2 * self.batch_size), 2 * self.batch_size, k=self.batch_size) 27 | mask = torch.from_numpy((diag + l1 + l2)) 28 | mask = (1 - mask).type(torch.bool) 29 | return mask.to(self.device) 30 | 31 | @staticmethod 32 | def _dot_simililarity(x, y): 33 | v = torch.tensordot(x.unsqueeze(1), y.T.unsqueeze(0), dims=2) 34 | # x shape: (N, 1, C) 35 | # y shape: (1, C, 2N) 36 | # v shape: (N, 2N) 37 | return v 38 | 39 | def _cosine_simililarity(self, x, y): 40 | # x shape: (N, 1, C) 41 | # y shape: (1, 2N, C) 42 | # v shape: (N, 2N) 43 | v = self._cosine_similarity(x.unsqueeze(1), y.unsqueeze(0)) 44 | return v 45 | 46 | def forward(self, zis, zjs): 47 | representations = torch.cat([zjs, zis], dim=0) 48 | 49 | similarity_matrix = self.similarity_function(representations, representations) 50 | 51 | # filter out the scores from the positive samples 52 | l_pos = torch.diag(similarity_matrix, self.batch_size) 53 | r_pos = torch.diag(similarity_matrix, -self.batch_size) 54 | positives = torch.cat([l_pos, r_pos]).view(2 * self.batch_size, 1) 55 | 56 | negatives = similarity_matrix[self.mask_samples_from_same_repr].view(2 * self.batch_size, -1) 57 | 58 | logits = torch.cat((positives, negatives), dim=1) 59 | logits /= self.temperature 60 | 61 | labels = torch.zeros(2 * self.batch_size).to(self.device).long() 62 | loss = self.criterion(logits, labels) 63 | 64 | return loss / (2 * self.batch_size) 65 | -------------------------------------------------------------------------------- /pc_test.py: -------------------------------------------------------------------------------- 1 | # from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 2 | from tools.test_dataloader import TestDataloader 3 | import numpy as np 4 | import os 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from torch.autograd import Variable 9 | import torchvision.models as models 10 | import argparse 11 | import torch.optim as optim 12 | import time 13 | import torchvision.models as models 14 | # from corrnet import SingleViewNet 15 | from models.dgcnn import get_graph_feature 16 | from models.dgcnn import DGCNN 17 | from models.pointnet_part_seg import PointnetPartSeg 18 | # from models.pointnet_part_seg import PointNet_Part 19 | from models.meshnet import MeshNet 20 | from models.SVCNN import Semi3D, SingleViewNet 21 | from tools.triplet_dataloader import TripletDataloader 22 | from tqdm import tqdm 23 | from tools.utils import calculate_accuracy 24 | 25 | 26 | 27 | def extract(dgcnn, num_views, split, exp_name): 28 | 29 | dataset = 'ModelNet40' 30 | 31 | # test_data_set = TripletDataloader(dataset = 'ModelNet40', num_points = 1024, partition='test', perceptange = 10) 32 | test_data_set = TestDataloader(dataset, num_points = 1024, num_views = 2, partition= 'test') 33 | test_data_loader = torch.utils.data.DataLoader(test_data_set, batch_size= 1, shuffle = False, num_workers=1) 34 | 35 | print('length of the dataset: ', len(test_data_set)) 36 | 37 | start_time = time.time() 38 | 39 | img_pred_lst = [] 40 | pt_pred_lst = [] 41 | mesh_pred_lst = [] 42 | fused_pred_lst = [] 43 | target_lst = [] 44 | 45 | iteration = 0 46 | for data in test_data_loader: 47 | 48 | pt, img, imgV, centers, corners, normals, neighbor_index, target = data #the last one is the target 49 | 50 | pt = Variable(pt).to('cuda') 51 | 52 | img = Variable(img).to('cuda') 53 | imgV = Variable(imgV).to('cuda') 54 | 55 | centers = Variable(torch.cuda.FloatTensor(centers.cuda())) 56 | corners = Variable(torch.cuda.FloatTensor(corners.cuda())) 57 | normals = Variable(torch.cuda.FloatTensor(normals.cuda())) 58 | neighbor_index = Variable(torch.cuda.LongTensor(neighbor_index.cuda())) 59 | 60 | target = torch.squeeze(target, dim = 1) 61 | target = Variable(target).to('cuda') 62 | 63 | pt = pt.permute(0,2,1) 64 | 65 | pt_pred, pt_feat, pt_base = dgcnn(pt) 66 | # img_pred, img_feat, img_base = img_net(img, imgV) 67 | # mesh_pred, mesh_feat, mesh_base = mesh_net(centers, corners, normals, neighbor_index) 68 | 69 | # fused_pred = mesh_pred * pt_pred 70 | 71 | 72 | pt_pred_lst.append(pt_pred.data.cpu()) 73 | # mesh_pred_lst.append(mesh_pred.data.cpu()) 74 | # img_pred_lst.append(img_pred.data.cpu()) 75 | # fused_pred_lst.append(fused_pred.data.cpu()) 76 | # print(target.size()) 77 | target_lst.append(target.data.cpu()) 78 | 79 | iteration = iteration + 1 80 | if iteration % 400 == 0: 81 | print('iterations: ', iteration) 82 | 83 | pt_pred = torch.cat(pt_pred_lst, dim = 0) 84 | # mesh_pred = torch.cat(mesh_pred_lst, dim = 0) 85 | # img_pred = torch.cat(img_pred_lst, dim = 0) 86 | # fused_pred = torch.cat(fused_pred_lst, dim = 0) 87 | target = torch.cat(target_lst, dim = 0) 88 | 89 | print('pred size: ', pt_pred.size(), target.size()) 90 | pt_acc = calculate_accuracy(pt_pred, target) 91 | 92 | print('the pt acc: ', pt_acc) 93 | 94 | 95 | 96 | def extract_features(args): 97 | 98 | iterations = 6000 99 | num_views = 2 # 1 12 80 100 | 101 | # weights_folder = 'ModelNet40-pt1024-5percent-supervised' 102 | weights_folder = 'ModelNet40-pt1024-5percent-entropy-Icontrast' 103 | 104 | dgcnn = DGCNN(args) 105 | dgcnn = torch.nn.DataParallel(dgcnn) 106 | dgcnn_name = './checkpoints/%s/%d-pt_net.pkl'%(weights_folder, iterations) 107 | dgcnn.load_state_dict(torch.load(dgcnn_name)['state_dict']) 108 | 109 | 110 | # img_net = img_net.eval() 111 | dgcnn = dgcnn.eval() 112 | # mesh_net = mesh_net.eval() 113 | # img_net = img_net.to('cuda') 114 | 115 | dgcnn = dgcnn.to('cuda') 116 | # mesh_net = mesh_net.to('cuda') 117 | 118 | 119 | print('evaluation for the testing split') 120 | extract(dgcnn, num_views, 'test', exp_name = weights_folder) 121 | print('------------------ Al the Features are saved ---------------------------') 122 | 123 | if __name__ == "__main__": 124 | # Training settings 125 | parser = argparse.ArgumentParser(description='RGB and Point Cloud Correspondence') 126 | 127 | parser.add_argument('--no_cuda', type=bool, default=False, 128 | help='enables CUDA training') 129 | parser.add_argument('--seed', type=int, default=1, metavar='S', 130 | help='random seed (default: 1)') 131 | 132 | parser.add_argument('--num_points', type=int, default=1024, 133 | help='num of points to use') 134 | 135 | parser.add_argument('--dropout', type=float, default=0.5, 136 | help='dropout rate') 137 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 138 | help='Dimension of embeddings') 139 | 140 | parser.add_argument('--k', type=int, default=20, metavar='N', 141 | help='Num of nearest neighbors to use') 142 | parser.add_argument('--model_path', type=str, default='', metavar='N', 143 | help='Pretrained model path') 144 | 145 | parser.add_argument('--gpu_id', type=str, default='0', 146 | help='GPU used to train the network') 147 | 148 | args = parser.parse_args() 149 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 150 | extract_features(args) -------------------------------------------------------------------------------- /pc_train.py: -------------------------------------------------------------------------------- 1 | from __future__ import division, absolute_import 2 | from models.dgcnn import DGCNN 3 | from models.pointnet_part_seg import PointnetPartSeg 4 | # from models.pointnet_part_seg import PointNet_Part 5 | from models.meshnet import MeshNet 6 | from models.SVCNN import Semi3D, SingleViewNet 7 | from tools.triplet_dataloader import TripletDataloader 8 | from tools.utils import calculate_accuracy 9 | import numpy as np 10 | import math 11 | import os 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as F 15 | from torch.autograd import Variable 16 | import torchvision.models as models 17 | import argparse 18 | import torch.optim as optim 19 | import time 20 | 21 | from nt_xent import NTXentLoss 22 | 23 | from torch.utils.tensorboard import SummaryWriter 24 | 25 | import warnings 26 | warnings.filterwarnings('ignore',category=FutureWarning) 27 | 28 | 29 | 30 | def training(args): 31 | if not os.path.exists(args.save): 32 | os.makedirs(args.save) 33 | 34 | # img_net = SingleViewNet(pre_trained = True) 35 | pt_net = DGCNN(args) 36 | # pt_net = PointnetPartSeg() 37 | 38 | # meshnet = MeshNet() 39 | 40 | # model = Semi3D(img_net, pt_net, meshnet) 41 | 42 | pt_net = pt_net.to('cuda') 43 | 44 | pt_net = torch.nn.DataParallel(pt_net) 45 | 46 | pt_net.train(True) 47 | 48 | optimizer = optim.SGD(pt_net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) 49 | 50 | writer = SummaryWriter(os.path.join(args.save, 'summary')) 51 | 52 | #data splittted into unlabeled/labeled/test 53 | labeled_set = TripletDataloader(dataset = 'ModelNet40', num_points = args.num_points, partition='labeled', perceptange = 5) 54 | labeled_data_loader = torch.utils.data.DataLoader(labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=8, drop_last=True) 55 | 56 | 57 | unlabeled_set = TripletDataloader(dataset = 'ModelNet40', num_points = args.num_points, partition='unlabeled', perceptange = 5) 58 | unlabeled_data_loader = torch.utils.data.DataLoader(unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=8, drop_last=True) 59 | 60 | print('************************************************************') 61 | print(' check the following important parametes ') 62 | print('the number of labeled sample: ', len(labeled_set)) 63 | # print('the number of unlabeled sample: ', len(unlabeled_set)) 64 | print('the temperature for the probability: ', args.T) 65 | print('the threshold for the probability: ', args.threshold) 66 | print('************************************************************') 67 | 68 | # The loss introduced in Hinton's paper 69 | nt_xent_criterion = NTXentLoss('cuda', args.batch_size, temperature = 0.5, use_cosine_similarity = True) 70 | # mse_criterion = nn.MSELoss() 71 | ce_criterion = nn.CrossEntropyLoss(reduction='mean') 72 | 73 | iteration = 0 74 | start_time = time.time() 75 | for epoch in range(args.epochs): 76 | for l_data, u_data in zip(labeled_data_loader, unlabeled_data_loader): 77 | pt1, pt1A, img1, img1V, centers1, corners1, normals1, neighbor_index1, target1 = l_data #the last one is the target 78 | pt2, pt2A, img2, img2V, centers2, corners2, normals2, neighbor_index2, target2 = u_data #the last one is the target 79 | 80 | pt1 = Variable(pt1).to('cuda') 81 | pt1 = pt1.permute(0,2,1) 82 | 83 | # pt1A = Variable(pt1A).to('cuda') 84 | # pt1A = pt1A.permute(0,2,1) 85 | 86 | pt2 = Variable(pt2).to('cuda') 87 | pt2 = pt2.permute(0,2,1) 88 | 89 | pt2A = Variable(pt2A).to('cuda') 90 | pt2A = pt2A.permute(0,2,1) 91 | 92 | target1 = torch.squeeze(target1) 93 | target1 = Variable(target1).to('cuda') 94 | 95 | optimizer.zero_grad() 96 | # 97 | # pt_pred1, mesh_pred1, img_pred1, fused_pred1, pt_feat1, mesh_feat1, img_feat1 = model(pt1, img1, img1V, centers1, corners1, normals1, neighbor_index1) 98 | pt_pred1, pt_feat1, pt_base1 = pt_net(pt1) 99 | # pt_pred1A, pt_feat1A, pt_base1A = pt_net(pt1A) 100 | 101 | pt_pred2, pt_feat2, pt_base2 = pt_net(pt2) 102 | pt_pred2A, pt_feat2A, pt_base2A = pt_net(pt2A) 103 | # pt_pred2, mesh_pred2, img_pred2, fused_pred2, pt_feat2, mesh_feat2, img_feat2 = model(pt2, img2, img2V, centers2, corners2, normals2, neighbor_index2) 104 | 105 | #cross-entropy loss on the labeled data 106 | pt_ce_loss = ce_criterion(pt_pred1, target1) 107 | 108 | 109 | #contrastive loss on the unlabeled dataset 110 | pt_contrast_loss = nt_xent_criterion(pt_feat2, pt_feat2A) 111 | 112 | loss = pt_ce_loss + pt_contrast_loss 113 | 114 | loss.backward() 115 | 116 | #update the parameters for the center_loss 117 | optimizer.step() 118 | 119 | pt_acc1 = calculate_accuracy(pt_pred1, target1) 120 | 121 | #classification accuracy on the labeld sample 122 | writer.add_scalar('pt_acc', pt_acc1, iteration) 123 | 124 | #Xentropy loss on the labeled data 125 | writer.add_scalar('loss/pt_ce_loss', pt_ce_loss.item(), iteration) 126 | writer.add_scalar('loss/pt_contrast_loss', pt_contrast_loss.item(), iteration) 127 | 128 | writer.add_scalar('loss/loss', loss.item(), iteration) 129 | 130 | 131 | if (iteration%args.lr_step) == 0: 132 | lr = args.lr * (0.1 ** (iteration // args.lr_step)) 133 | print('New LR: ' + str(lr)) 134 | for param_group in optimizer.param_groups: 135 | param_group['lr'] = lr 136 | 137 | if iteration % args.per_print == 0: 138 | print('[%d][%d] loss: %.2f acc: %.2f time: %.2f vid: %d' % \ 139 | (epoch, iteration, loss.item(), pt_acc1, time.time() - start_time, 2 * pt1.size(0))) 140 | start_time = time.time() 141 | 142 | iteration = iteration + 1 143 | if((iteration+1) % args.per_save) ==0: 144 | print('----------------- Save The Network ------------------------') 145 | 146 | pt_net_name = args.save + str(iteration+1)+'-pt_net.pkl' 147 | torch.save({'state_dict': pt_net.state_dict()}, pt_net_name) 148 | 149 | iteration = iteration + 1 150 | if iteration > args.max_step: 151 | return 152 | 153 | 154 | if __name__ == "__main__": 155 | # Training settings 156 | parser = argparse.ArgumentParser(description='Learning View and Model invariant features for 3D shapes') 157 | 158 | parser.add_argument('--batch_size', type=int, default=42, metavar='batch_size', 159 | help='Size of batch)') 160 | 161 | parser.add_argument('--epochs', type=int, default=2000, metavar='N', 162 | help='number of episode to train ') 163 | #optimizer 164 | parser.add_argument('--lr', type=float, default=0.005, metavar='LR', 165 | help='learning rate (default: 0.001, 0.1 if using sgd)') 166 | 167 | parser.add_argument('--lr_step', type=int, default = 2000, 168 | help='how many iterations to decrease the learning rate') 169 | 170 | parser.add_argument('--max_step', type=int, default = 6100, 171 | help='maximum steps to train the network') 172 | 173 | parser.add_argument('--momentum', type=float, default=0.9, metavar='M', 174 | help='SGD momentum (default: 0.9)') 175 | 176 | parser.add_argument('--no_cuda', type=bool, default=False, 177 | help='enables CUDA training') 178 | 179 | parser.add_argument('--seed', type=int, default=1, metavar='S', 180 | help='random seed (default: 1)') 181 | 182 | parser.add_argument('--T', type=int, default = 1, 183 | help='temperature for the prediction') 184 | 185 | parser.add_argument('--threshold', type=int, default = 0.95, 186 | help='threshold for the positive samples') 187 | 188 | #image for SVCNN 189 | parser.add_argument('--num_views', type=int, default=180, metavar='S', 190 | help='number of views for training (default: 6)') 191 | #DGCNN 192 | parser.add_argument('--num_points', type=int, default=1024, 193 | help='num of points to use') 194 | parser.add_argument('--dropout', type=float, default=0.5, 195 | help='dropout rate') 196 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 197 | help='Dimension of embeddings') 198 | parser.add_argument('--k', type=int, default=20, metavar='N', 199 | help='Num of nearest neighbors to use') 200 | 201 | parser.add_argument('--model_path', type=str, default='', metavar='N', 202 | help='Pretrained model path') 203 | 204 | parser.add_argument('--weight_decay', type=float, default=1e-3, metavar='weight_decay', 205 | help='learning rate (default: 1e-3)') 206 | 207 | parser.add_argument('--per_save', type=int, default=2000, 208 | help='how many iterations to save the model') 209 | 210 | parser.add_argument('--per_print', type=int, default=100, 211 | help='how many iterations to print the loss and accuracy') 212 | parser.add_argument('--save', type=str, default='./checkpoints/ModelNet40-pt1024-5percent-entropy-Icontrast/', 213 | help='path to save the final model') 214 | 215 | parser.add_argument('--gpu_id', type=str, default='0,1,2', 216 | help='GPU used to train the network') 217 | 218 | parser.add_argument('--log', type=str, default='log/', 219 | help='path to the log information') 220 | 221 | args = parser.parse_args() 222 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 223 | torch.backends.cudnn.enabled = False 224 | training(args) 225 | -------------------------------------------------------------------------------- /test_point.py: -------------------------------------------------------------------------------- 1 | # from tools.dual_dataloader import SingleViewDataloader, MultiViewDataloader 2 | from tools.test_dataloader import TestDataloader 3 | import numpy as np 4 | import os 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | from torch.autograd import Variable 9 | import torchvision.models as models 10 | import argparse 11 | import torch.optim as optim 12 | import time 13 | import torchvision.models as models 14 | # from corrnet import SingleViewNet 15 | from models.dgcnn import get_graph_feature 16 | from models.dgcnn import DGCNN 17 | from models.pointnet_part_seg import PointnetPartSeg 18 | # from models.pointnet_part_seg import PointNet_Part 19 | from models.meshnet import MeshNet 20 | from models.SVCNN import Semi3D, SingleViewNet, FusionNet 21 | from tools.triplet_dataloader import TripletDataloader 22 | from tqdm import tqdm 23 | from tools.utils import calculate_accuracy 24 | 25 | 26 | 27 | def extract( dgcnn, num_views, split, exp_name): 28 | 29 | dataset = 'ModelNet40' 30 | 31 | # test_data_set = TripletDataloader(dataset = 'ModelNet40', num_points = 1024, partition='test', perceptange = 10) 32 | test_data_set = TestDataloader(dataset, num_points = 1024, num_views = 2, partition= 'test') 33 | test_data_loader = torch.utils.data.DataLoader(test_data_set, batch_size= 1, shuffle = False, num_workers=1) 34 | 35 | print('length of the dataset: ', len(test_data_set)) 36 | 37 | start_time = time.time() 38 | 39 | img_pred_lst = [] 40 | pt_pred_lst = [] 41 | mesh_pred_lst = [] 42 | fused_pred_lst = [] 43 | mean_pred_lst = [] 44 | geometric_mean_pred_lst = [] 45 | target_lst = [] 46 | 47 | iteration = 0 48 | for data in test_data_loader: 49 | 50 | pt, img, imgV, centers, corners, normals, neighbor_index, target = data #the last one is the target 51 | 52 | pt = Variable(pt).to('cuda') 53 | 54 | img = Variable(img).to('cuda') 55 | imgV = Variable(imgV).to('cuda') 56 | 57 | centers = Variable(torch.cuda.FloatTensor(centers.cuda())) 58 | corners = Variable(torch.cuda.FloatTensor(corners.cuda())) 59 | normals = Variable(torch.cuda.FloatTensor(normals.cuda())) 60 | neighbor_index = Variable(torch.cuda.LongTensor(neighbor_index.cuda())) 61 | 62 | target = torch.squeeze(target, dim = 1) 63 | target = Variable(target).to('cuda') 64 | 65 | pt = pt.permute(0,2,1) 66 | 67 | pt_pred, pt_feat, pt_base = dgcnn(pt) 68 | 69 | 70 | 71 | pt_pred_lst.append(pt_pred.data.cpu()) 72 | 73 | 74 | # print(target.size()) 75 | target_lst.append(target.data.cpu()) 76 | 77 | iteration = iteration + 1 78 | if iteration % 400 == 0: 79 | print('iterations: ', iteration) 80 | 81 | pt_pred = torch.cat(pt_pred_lst, dim = 0) 82 | 83 | target = torch.cat(target_lst, dim = 0) 84 | 85 | # print('pred size: ', img_pred.size(), mesh_pred.size(), target.size()) 86 | pt_acc = calculate_accuracy(pt_pred, target) 87 | 88 | 89 | 90 | print('the pt acc: %.4f'%(pt_acc)) 91 | 92 | 93 | # img_acc = calculate_accuracy(img_pred, target) 94 | # pt_acc = calculate_accuracy(pt_pred, target) 95 | # mesh_acc = calculate_accuracy(mesh_pred, target) 96 | # fused_acc = calculate_accuracy(fused_pred, target) 97 | # # print(img_acc, pt_acc, mesh_acc) 98 | 99 | # img_acc_lst.append(img_acc) 100 | # pt_acc_lst.append(pt_acc) 101 | # mesh_acc_lst.append(mesh_acc) 102 | # fused_acc_lst.append(fused_acc) 103 | 104 | # iteration = iteration + 1 105 | # if iteration % 200 == 0: 106 | # print('the pt acc: ', sum(pt_acc_lst)/len(pt_acc_lst)) 107 | # print('the img acc: ', sum(img_acc_lst)/len(img_acc_lst)) 108 | # print('the mesh acc: ', sum(mesh_acc_lst)/len(mesh_acc_lst)) 109 | # print('the fused acc: ', sum(fused_acc_lst)/len(fused_acc_lst)) 110 | # print('-------------------------------------------------') 111 | 112 | # print('the pt acc: ', sum(pt_acc_lst)/len(pt_acc_lst)) 113 | # print('the img acc: ', sum(img_acc_lst)/len(img_acc_lst)) 114 | # print('the mesh acc: ', sum(mesh_acc_lst)/len(mesh_acc_lst)) 115 | # print('the fused acc: ', sum(fused_acc_lst)/len(fused_acc_lst)) 116 | 117 | 118 | def extract_features(args): 119 | 120 | iterations = 10000 121 | num_views = 2 # 1 12 80 122 | 123 | # weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-Xcontrast-MultiAgreement-T095-Fused-Warmup-2percent' 124 | # weights_folder = 'ModelNet40-pt1024-mesh-img56-Xentropy-2percent-supervised' 125 | weights_folder = 'ModelNet40-pt1024-5percent-entropy-Icontrast_task0_p2' 126 | 127 | # img_net = torch.nn.DataParallel(img_net) 128 | 129 | 130 | dgcnn = DGCNN(args) 131 | dgcnn = torch.nn.DataParallel(dgcnn) 132 | dgcnn_name = './checkpoints/%s/%d-pt_net.pkl'%(weights_folder, iterations) 133 | dgcnn.load_state_dict(torch.load(dgcnn_name)['state_dict']) 134 | 135 | 136 | 137 | 138 | dgcnn = dgcnn.eval() 139 | 140 | 141 | dgcnn = dgcnn.to('cuda') 142 | 143 | 144 | 145 | print('evaluation for the testing split') 146 | extract( dgcnn, num_views, 'test', exp_name = weights_folder) 147 | print('------------------ Al the Features are saved ---------------------------') 148 | 149 | if __name__ == "__main__": 150 | # Training settings 151 | parser = argparse.ArgumentParser(description='RGB and Point Cloud Correspondence') 152 | 153 | parser.add_argument('--no_cuda', type=bool, default=False, 154 | help='enables CUDA training') 155 | parser.add_argument('--seed', type=int, default=1, metavar='S', 156 | help='random seed (default: 1)') 157 | 158 | parser.add_argument('--num_points', type=int, default=1024, 159 | help='num of points to use') 160 | 161 | parser.add_argument('--dropout', type=float, default=0.5, 162 | help='dropout rate') 163 | parser.add_argument('--emb_dims', type=int, default=1024, metavar='N', 164 | help='Dimension of embeddings') 165 | 166 | parser.add_argument('--k', type=int, default=20, metavar='N', 167 | help='Num of nearest neighbors to use') 168 | parser.add_argument('--model_path', type=str, default='', metavar='N', 169 | help='Pretrained model path') 170 | 171 | parser.add_argument('--gpu_id', type=str, default='0', 172 | help='GPU used to train the network') 173 | 174 | args = parser.parse_args() 175 | os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id 176 | extract_features(args) -------------------------------------------------------------------------------- /tools/.ipynb_checkpoints/test_dataloader-checkpoint.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import glob 4 | import h5py 5 | import json 6 | import numpy as np 7 | import random 8 | from PIL import Image 9 | from torch.utils.data import Dataset 10 | from torchvision import transforms, datasets 11 | import torch 12 | from tools.visualize import showpoints 13 | from tools.dual_dataloader import load_data, load_modelnet10_data 14 | 15 | 16 | class TestDataloader(Dataset): 17 | def __init__(self, dataset, num_points, num_views, partition='test'): 18 | self.dataset = dataset 19 | if self.dataset == 'ModelNet40': 20 | self.data, self.label, self.img_lst = load_data(partition) 21 | else: 22 | self.data, self.label, self.img_lst = load_modelnet10_data(partition) 23 | 24 | self.num_points = num_points 25 | self.num_views = num_views 26 | self.partition = partition 27 | 28 | self.img_transform = transforms.Compose([ 29 | # transforms.CenterCrop(224), 30 | transforms.Resize(224), 31 | # transforms.Resize((450, 800)), 32 | transforms.ToTensor(), 33 | transforms.Normalize(mean=[0.485, 0.456, 0.406], 34 | std=[0.229, 0.224, 0.225]) 35 | ]) 36 | 37 | def get_data(self, item): 38 | 39 | # Get Image Data first 40 | names = self.img_lst[item] 41 | names = names.split('/') 42 | 43 | #random select one image from the 12 images for each object 44 | 45 | img_idx_lst = [x for x in range(180)] 46 | random.shuffle(img_idx_lst) 47 | img_idx = img_idx_lst[0] 48 | img_names = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx) 49 | img = Image.open(img_names).convert('RGB') 50 | img_idx2 = img_idx_lst[1] 51 | img_name2 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx2) 52 | img2 = Image.open(img_name2).convert('RGB') 53 | 54 | 55 | 56 | label = self.label[item] 57 | 58 | pointcloud = self.data[item] 59 | choice = np.random.choice(len(pointcloud), self.num_points, replace=True) 60 | pointcloud = pointcloud[choice, :] 61 | 62 | img = self.img_transform(img) 63 | img2 = self.img_transform(img2) 64 | 65 | return pointcloud, label, img, img2 66 | 67 | 68 | def check_exist(self, item): 69 | names = self.img_lst[item] 70 | names = names.split('/') 71 | # mesh_path = './dataset/ModelNet40_Mesh/' + names[0] + '/train/' + names[1][:-4] + '.npz' 72 | mesh_path = os.path.join('/scratch1/zhiminc/MVdata/data/ModelNet40_Mesh/', names[0], self.partition, '%s.npz'%(names[1][:-4])) 73 | 74 | return os.path.isfile(mesh_path) 75 | 76 | 77 | def get_mesh(self, item): 78 | names = self.img_lst[item] 79 | names = names.split('/') 80 | # mesh_path = './dataset/ModelNet40_Mesh/' + names[0] + '/train/' + names[1][:-4] + '.npz' 81 | mesh_path = os.path.join('/scratch1/zhiminc/MVdata/data/ModelNet40_Mesh/', names[0], self.partition, '%s.npz'%(names[1][:-4])) 82 | 83 | data = np.load(mesh_path) 84 | face = data['face'] 85 | neighbor_index = data['neighbor_index'] 86 | 87 | # no data augmentation in the testing dataloader 88 | # if self.partition == 'train': 89 | # sigma, clip = 0.01, 0.05 90 | # jittered_data = np.clip(sigma * np.random.randn(*face[:, :12].shape), -1 * clip, clip) 91 | # face = np.concatenate((face[:, :12] + jittered_data, face[:, 12:]), 1) 92 | 93 | # fill for n < max_faces with randomly picked faces 94 | max_faces = 1024 95 | num_point = len(face) 96 | if num_point < max_faces: 97 | fill_face = [] 98 | fill_neighbor_index = [] 99 | for i in range(max_faces - num_point): 100 | index = np.random.randint(0, num_point) 101 | fill_face.append(face[index]) 102 | fill_neighbor_index.append(neighbor_index[index]) 103 | face = np.concatenate((face, np.array(fill_face))) 104 | neighbor_index = np.concatenate((neighbor_index, np.array(fill_neighbor_index))) 105 | 106 | # to tensor 107 | face = torch.from_numpy(face).float() 108 | neighbor_index = torch.from_numpy(neighbor_index).long() 109 | 110 | # reorganize 111 | face = face.permute(1, 0).contiguous() 112 | centers, corners, normals = face[:3], face[3:12], face[12:] 113 | corners = corners - torch.cat([centers, centers, centers], 0) 114 | 115 | return centers, corners, normals, neighbor_index 116 | 117 | def __getitem__(self, item): 118 | #name of this point cloud 119 | 120 | while not self.check_exist(item): 121 | idx = random.randint(0, len(self.data)-1) 122 | item = idx 123 | 124 | pt, target, img, img_v = self.get_data(item) 125 | pt = torch.from_numpy(pt) 126 | 127 | centers, corners, normals, neighbor_index = self.get_mesh(item) 128 | 129 | # return pt1, img1, img2, label1 130 | return pt, img, img_v, centers, corners, normals, neighbor_index, target 131 | 132 | def __len__(self): 133 | return self.data.shape[0] 134 | 135 | if __name__ == '__main__': 136 | 137 | train_set = SingleViewDataloader(num_points = 1024, partition='train') 138 | 139 | data_loader_loader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=False,num_workers=1) 140 | 141 | cnt = 0 142 | for data in data_loader_loader: 143 | pt1, pt2, img1, img2, img1_v, img2_v, label1, label2, pos, neg = data 144 | pt1 = pt1.numpy() 145 | pt1 = pt1[0,:,:] 146 | # print(pt1.shape) 147 | # print(pt1) 148 | # print(np.amin(pt1), np.amax(pt1)) 149 | # showpoints(pt1) 150 | -------------------------------------------------------------------------------- /tools/.ipynb_checkpoints/test_dataloader_twfview-checkpoint.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import glob 4 | import h5py 5 | import json 6 | import numpy as np 7 | import random 8 | from PIL import Image 9 | from torch.utils.data import Dataset 10 | from torchvision import transforms, datasets 11 | import torch 12 | from tools.visualize import showpoints 13 | from tools.dual_dataloader import load_data, load_modelnet10_data 14 | 15 | 16 | class TestDataloader(Dataset): 17 | def __init__(self, dataset, num_points, num_views, partition='test'): 18 | self.dataset = dataset 19 | if self.dataset == 'ModelNet40': 20 | self.data, self.label, self.img_lst = load_data(partition) 21 | else: 22 | self.data, self.label, self.img_lst = load_modelnet10_data(partition) 23 | 24 | self.num_points = num_points 25 | self.num_views = num_views 26 | self.partition = partition 27 | 28 | self.img_transform = transforms.Compose([ 29 | # transforms.CenterCrop(224), 30 | transforms.Resize(224), 31 | # transforms.Resize((450, 800)), 32 | transforms.ToTensor(), 33 | transforms.Normalize(mean=[0.485, 0.456, 0.406], 34 | std=[0.229, 0.224, 0.225]) 35 | ]) 36 | 37 | def get_data(self, item): 38 | 39 | # Get Image Data first 40 | names = self.img_lst[item] 41 | names = names.split('/') 42 | 43 | #random select one image from the 12 images for each object 44 | 45 | img_idx_lst = [x for x in range(180)] 46 | random.shuffle(img_idx_lst) 47 | img_idx = img_idx_lst[0] 48 | img_names = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx) 49 | img = Image.open(img_names).convert('RGB') 50 | 51 | img_idx2 = img_idx_lst[1] 52 | img_name2 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx2) 53 | img2 = Image.open(img_name2).convert('RGB') 54 | 55 | img_idx3 = img_idx_lst[2] 56 | img_name3 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx3) 57 | img3 = Image.open(img_name3).convert('RGB') 58 | 59 | img_idx4 = img_idx_lst[3] 60 | img_name4 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx4) 61 | img4 = Image.open(img_name4).convert('RGB') 62 | 63 | img_idx5 = img_idx_lst[4] 64 | img_name5 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx5) 65 | img5 = Image.open(img_name5).convert('RGB') 66 | 67 | img_idx6 = img_idx_lst[5] 68 | img_name6 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx6) 69 | img6 = Image.open(img_name6).convert('RGB') 70 | 71 | img_idx7 = img_idx_lst[6] 72 | img_name7 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx7) 73 | img7 = Image.open(img_name7).convert('RGB') 74 | 75 | img_idx8 = img_idx_lst[7] 76 | img_name8 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx8) 77 | img8 = Image.open(img_name8).convert('RGB') 78 | 79 | img_idx9 = img_idx_lst[8] 80 | img_name9 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx9) 81 | img9 = Image.open(img_name9).convert('RGB') 82 | 83 | img_idx10 = img_idx_lst[9] 84 | img_name10 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx10) 85 | img10 = Image.open(img_name10).convert('RGB') 86 | 87 | img_idx11 = img_idx_lst[10] 88 | img_name11 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx11) 89 | img11 = Image.open(img_name11).convert('RGB') 90 | 91 | img_idx12 = img_idx_lst[11] 92 | img_name12 = '/scratch1/zhiminc/MVdata/data/%s-Images-180/%s/%s/%s.%d.png' % (self.dataset, names[0], names[1][:-4], names[1][:-4], img_idx12) 93 | img12 = Image.open(img_name12).convert('RGB') 94 | 95 | 96 | 97 | label = self.label[item] 98 | 99 | pointcloud = self.data[item] 100 | choice = np.random.choice(len(pointcloud), self.num_points, replace=True) 101 | pointcloud = pointcloud[choice, :] 102 | 103 | img = self.img_transform(img) 104 | img2 = self.img_transform(img2) 105 | img3 = self.img_transform(img3) 106 | img4 = self.img_transform(img4) 107 | img5 = self.img_transform(img5) 108 | img6 = self.img_transform(img6) 109 | img7 = self.img_transform(img7) 110 | img8 = self.img_transform(img8) 111 | img9 = self.img_transform(img9) 112 | img10 = self.img_transform(img10) 113 | img11 = self.img_transform(img11) 114 | img12 = self.img_transform(img12) 115 | 116 | return pointcloud, label, img, img2, img3, img4, img5, img6, img7, img8, img9, img10, img11,img12 117 | 118 | 119 | def check_exist(self, item): 120 | names = self.img_lst[item] 121 | names = names.split('/') 122 | # mesh_path = './dataset/ModelNet40_Mesh/' + names[0] + '/train/' + names[1][:-4] + '.npz' 123 | mesh_path = os.path.join('/scratch1/zhiminc/MVdata/data/ModelNet40_Mesh/', names[0], self.partition, '%s.npz'%(names[1][:-4])) 124 | 125 | return os.path.isfile(mesh_path) 126 | 127 | 128 | def get_mesh(self, item): 129 | names = self.img_lst[item] 130 | names = names.split('/') 131 | # mesh_path = './dataset/ModelNet40_Mesh/' + names[0] + '/train/' + names[1][:-4] + '.npz' 132 | mesh_path = os.path.join('/scratch1/zhiminc/MVdata/data/ModelNet40_Mesh/', names[0], self.partition, '%s.npz'%(names[1][:-4])) 133 | 134 | data = np.load(mesh_path) 135 | face = data['face'] 136 | neighbor_index = data['neighbor_index'] 137 | 138 | # no data augmentation in the testing dataloader 139 | # if self.partition == 'train': 140 | # sigma, clip = 0.01, 0.05 141 | # jittered_data = np.clip(sigma * np.random.randn(*face[:, :12].shape), -1 * clip, clip) 142 | # face = np.concatenate((face[:, :12] + jittered_data, face[:, 12:]), 1) 143 | 144 | # fill for n < max_faces with randomly picked faces 145 | max_faces = 1024 146 | num_point = len(face) 147 | if num_point < max_faces: 148 | fill_face = [] 149 | fill_neighbor_index = [] 150 | for i in range(max_faces - num_point): 151 | index = np.random.randint(0, num_point) 152 | fill_face.append(face[index]) 153 | fill_neighbor_index.append(neighbor_index[index]) 154 | face = np.concatenate((face, np.array(fill_face))) 155 | neighbor_index = np.concatenate((neighbor_index, np.array(fill_neighbor_index))) 156 | 157 | # to tensor 158 | face = torch.from_numpy(face).float() 159 | neighbor_index = torch.from_numpy(neighbor_index).long() 160 | 161 | # reorganize 162 | face = face.permute(1, 0).contiguous() 163 | centers, corners, normals = face[:3], face[3:12], face[12:] 164 | corners = corners - torch.cat([centers, centers, centers], 0) 165 | 166 | return centers, corners, normals, neighbor_index 167 | 168 | def __getitem__(self, item): 169 | #name of this point cloud 170 | 171 | while not self.check_exist(item): 172 | idx = random.randint(0, len(self.data)-1) 173 | item = idx 174 | 175 | pt, target, img, img_v, img3, img4, img5, img6, img7, img8, img9, img10, img11,img12 = self.get_data(item) 176 | pt = torch.from_numpy(pt) 177 | 178 | centers, corners, normals, neighbor_index = self.get_mesh(item) 179 | 180 | # return pt1, img1, img2, label1 181 | return pt, img, img_v,img3, img4, img5, img6, img7, img8, img9, img10, img11,img12, centers, corners, normals, neighbor_index, target 182 | 183 | def __len__(self): 184 | return self.data.shape[0] 185 | 186 | if __name__ == '__main__': 187 | 188 | train_set = SingleViewDataloader(num_points = 1024, partition='train') 189 | 190 | data_loader_loader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=False,num_workers=1) 191 | 192 | cnt = 0 193 | for data in data_loader_loader: 194 | pt1, pt2, img1, img2, img1_v, img2_v, label1, label2, pos, neg = data 195 | pt1 = pt1.numpy() 196 | pt1 = pt1[0,:,:] 197 | # print(pt1.shape) 198 | # print(pt1) 199 | # print(np.amin(pt1), np.amax(pt1)) 200 | # showpoints(pt1) 201 | -------------------------------------------------------------------------------- /tools/.ipynb_checkpoints/visualize-checkpoint.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import ctypes as ct 3 | import cv2 4 | import sys 5 | showsz = 800 6 | mousex, mousey = 0.5, 0.5 7 | zoom = 1.0 8 | changed = True 9 | 10 | def onmouse(*args): 11 | global mousex, mousey, changed 12 | y = args[1] 13 | x = args[2] 14 | mousex = x / float(showsz) 15 | mousey = y / float(showsz) 16 | changed = True 17 | 18 | # cv2.namedWindow('show3d') 19 | # cv2.moveWindow('show3d', 0, 0) 20 | # cv2.setMouseCallback('show3d', onmouse) 21 | 22 | dll = np.ctypeslib.load_library('tools/render_balls_so', '.') 23 | 24 | def showpoints(xyz, c_gt=None, c_pred = None, waittime=0, 25 | showrot=True, magnifyBlue=0, freezerot=False, background=(0,0,0), 26 | normalizecolor=True, ballradius=10): 27 | global showsz, mousex, mousey, zoom, changed 28 | xyz=xyz-xyz.mean(axis=0) 29 | radius=((xyz**2).sum(axis=-1)**0.5).max() 30 | xyz/=(radius*2.2)/showsz 31 | if c_gt is None: 32 | c0 = np.zeros((len(xyz), ), dtype='float32') + 255 33 | c1 = np.zeros((len(xyz), ), dtype='float32') + 255 34 | c2 = np.zeros((len(xyz), ), dtype='float32') + 255 35 | else: 36 | c0 = c_gt[:, 0] 37 | c1 = c_gt[:, 1] 38 | c2 = c_gt[:, 2] 39 | 40 | 41 | if normalizecolor: 42 | c0 /= (c0.max() + 1e-14) / 255.0 43 | c1 /= (c1.max() + 1e-14) / 255.0 44 | c2 /= (c2.max() + 1e-14) / 255.0 45 | 46 | 47 | c0 = np.require(c0, 'float32', 'C') 48 | c1 = np.require(c1, 'float32', 'C') 49 | c2 = np.require(c2, 'float32', 'C') 50 | 51 | show = np.zeros((showsz, showsz, 3), dtype='uint8') 52 | def render(): 53 | rotmat=np.eye(3) 54 | if not freezerot: 55 | xangle=(mousey-0.5)*np.pi*1.2 56 | else: 57 | xangle=0 58 | rotmat = rotmat.dot( 59 | np.array([ 60 | [1.0, 0.0, 0.0], 61 | [0.0, np.cos(xangle), -np.sin(xangle)], 62 | [0.0, np.sin(xangle), np.cos(xangle)], 63 | ])) 64 | if not freezerot: 65 | yangle = (mousex - 0.5) * np.pi * 1.2 66 | else: 67 | yangle = 0 68 | rotmat = rotmat.dot( 69 | np.array([ 70 | [np.cos(yangle), 0.0, -np.sin(yangle)], 71 | [0.0, 1.0, 0.0], 72 | [np.sin(yangle), 0.0, np.cos(yangle)], 73 | ])) 74 | rotmat *= zoom 75 | nxyz = xyz.dot(rotmat) + [showsz / 2, showsz / 2, 0] 76 | 77 | ixyz = nxyz.astype('int32') 78 | show[:] = background 79 | dll.render_ball( 80 | ct.c_int(show.shape[0]), ct.c_int(show.shape[1]), 81 | show.ctypes.data_as(ct.c_void_p), ct.c_int(ixyz.shape[0]), 82 | ixyz.ctypes.data_as(ct.c_void_p), c0.ctypes.data_as(ct.c_void_p), 83 | c1.ctypes.data_as(ct.c_void_p), c2.ctypes.data_as(ct.c_void_p), 84 | ct.c_int(ballradius)) 85 | 86 | if magnifyBlue > 0: 87 | show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( 88 | show[:, :, 0], 1, axis=0)) 89 | if magnifyBlue >= 2: 90 | show[:, :, 0] = np.maximum(show[:, :, 0], 91 | np.roll(show[:, :, 0], -1, axis=0)) 92 | show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( 93 | show[:, :, 0], 1, axis=1)) 94 | if magnifyBlue >= 2: 95 | show[:, :, 0] = np.maximum(show[:, :, 0], 96 | np.roll(show[:, :, 0], -1, axis=1)) 97 | if showrot: 98 | cv2.putText(show, 'xangle %d' % (int(xangle / np.pi * 180)), 99 | (30, showsz - 30), 0, 0.5, (255, 0, 0)) 100 | cv2.putText(show, 'yangle %d' % (int(yangle / np.pi * 180)), 101 | (30, showsz - 50), 0, 0.5, (255, 0, 0)) 102 | cv2.putText(show, 'zoom %d%%' % (int(zoom * 100)), (30, showsz - 70), 0, 103 | 0.5, (255, 0, 0)) 104 | changed = True 105 | while True: 106 | if changed: 107 | render() 108 | changed = False 109 | cv2.imshow('show3d', show) 110 | if waittime == 0: 111 | cmd = cv2.waitKey(10) % 256 112 | else: 113 | cmd = cv2.waitKey(waittime) % 256 114 | if cmd == ord('q'): 115 | break 116 | elif cmd == ord('Q'): 117 | sys.exit(0) 118 | 119 | if cmd == ord('t') or cmd == ord('p'): 120 | if cmd == ord('t'): 121 | if c_gt is None: 122 | c0 = np.zeros((len(xyz), ), dtype='float32') + 255 123 | c1 = np.zeros((len(xyz), ), dtype='float32') + 255 124 | c2 = np.zeros((len(xyz), ), dtype='float32') + 255 125 | else: 126 | c0 = c_gt[:, 0] 127 | c1 = c_gt[:, 1] 128 | c2 = c_gt[:, 2] 129 | else: 130 | if c_pred is None: 131 | c0 = np.zeros((len(xyz), ), dtype='float32') + 255 132 | c1 = np.zeros((len(xyz), ), dtype='float32') + 255 133 | c2 = np.zeros((len(xyz), ), dtype='float32') + 255 134 | else: 135 | c0 = c_pred[:, 0] 136 | c1 = c_pred[:, 1] 137 | c2 = c_pred[:, 2] 138 | if normalizecolor: 139 | c0 /= (c0.max() + 1e-14) / 255.0 140 | c1 /= (c1.max() + 1e-14) / 255.0 141 | c2 /= (c2.max() + 1e-14) / 255.0 142 | c0 = np.require(c0, 'float32', 'C') 143 | c1 = np.require(c1, 'float32', 'C') 144 | c2 = np.require(c2, 'float32', 'C') 145 | changed = True 146 | 147 | if cmd==ord('n'): 148 | zoom*=1.1 149 | changed=True 150 | elif cmd==ord('m'): 151 | zoom/=1.1 152 | changed=True 153 | elif cmd==ord('r'): 154 | zoom=1.0 155 | changed=True 156 | elif cmd==ord('s'): 157 | cv2.imwrite('show3d.png',show) 158 | if waittime!=0: 159 | break 160 | return cmd -------------------------------------------------------------------------------- /tools/.python-version: -------------------------------------------------------------------------------- 1 | anaconda3-5.0.0 2 | -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__init__.py -------------------------------------------------------------------------------- /tools/__pycache__/__init__.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/__init__.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/dual_dataloader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/dual_dataloader.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/dual_dataloader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/dual_dataloader.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/dual_dataloader.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/dual_dataloader.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader_twfview.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader_twfview.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader_twfview.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader_twfview.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/test_dataloader_twfview.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/test_dataloader_twfview.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/triplet_dataloader.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/triplet_dataloader.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/triplet_dataloader.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/triplet_dataloader.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/triplet_dataloader.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/triplet_dataloader.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/triplet_dataloader_twfview.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/triplet_dataloader_twfview.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/triplet_dataloader_twfview.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/triplet_dataloader_twfview.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/utils.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/utils.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/utils.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/utils.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /tools/__pycache__/visualize.cpython-36.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/visualize.cpython-36.pyc -------------------------------------------------------------------------------- /tools/__pycache__/visualize.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/visualize.cpython-37.pyc -------------------------------------------------------------------------------- /tools/__pycache__/visualize.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AutoAILab/M2CP-Learning/d15a88e643b20358c141e705dc6b7ed2c9a63f1e/tools/__pycache__/visualize.cpython-38.pyc -------------------------------------------------------------------------------- /tools/build.sh: -------------------------------------------------------------------------------- 1 | SCRIPT=`realpath $0` 2 | SCRIPTPATH=`dirname $SCRIPT` 3 | echo $SCRIPTPATH 4 | 5 | g++ -std=c++11 $SCRIPTPATH//render_balls_so.cpp -o $SCRIPTPATH/render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0 -------------------------------------------------------------------------------- /tools/gen_cloud_images.py: -------------------------------------------------------------------------------- 1 | import numpy as n 2 | import json 3 | import glob 4 | 5 | path_cloud = '../modelnet40_ply_hdf5_2048' 6 | # path_image = '../modelnet40_images_new_12x' 7 | path_image = '../modelnet40v2' 8 | 9 | all_cloud_files = sorted(glob.glob(path_cloud+'/*id2file.json')) 10 | # '../modelnet40_ply_hdf5_2048/ply_data_test_0_id2file.json' 11 | 12 | for f in all_cloud_files: 13 | with open(f) as json_file: 14 | 15 | set_ = f.split('_')[-3] 16 | num = f.split('_')[-2] 17 | print('set: ', set_, ' num: ', num) 18 | out = [] 19 | # out_path = path_cloud +'/ply_data_'+set_+num+ '.json' 20 | out_path = path_cloud +'/ply_data_'+set_+num+ '_80.json' 21 | print('name: ' + out_path) 22 | 23 | data = json.load(json_file) 24 | # print(data[0]) # bookshelf/bookshelf_0659.ply 25 | for path in data: 26 | clss = path.split('/')[0] 27 | name = path.split('/')[1][:-4] 28 | # tmp = name.split('_') 29 | idx = int(name[name.rfind('_')+1:]) 30 | print(path, name, idx, '%09d'%idx, name[:name.rfind('_')] + '_%09d'%idx) 31 | new_name = name[:name.rfind('_')] + '_%09d'%idx 32 | print(path_image+'/'+clss+'/'+set_+'/'+new_name+'*.jpg') 33 | images = sorted(glob.glob(path_image+'/'+clss+'/'+set_+'/'+new_name+'*.jpg')) 34 | images = ['/'.join(p.split('/')[-3:]) for p in images] 35 | print(len(images)) 36 | out.append(images) 37 | 38 | with open(out_path, 'w') as outfile: 39 | json.dump(out, outfile) -------------------------------------------------------------------------------- /tools/gen_cloud_images_shapenet.py: -------------------------------------------------------------------------------- 1 | import numpy as n 2 | import json 3 | import glob 4 | 5 | path_cloud = '../shapenetcorev2_hdf5_2048' 6 | path_image = '../shapenet55v2' 7 | 8 | all_cloud_files = sorted(glob.glob(path_cloud+'/*id2file.json')) 9 | # '../modelnet40_ply_hdf5_2048/test0_id2file.json' 10 | 11 | for f in all_cloud_files: 12 | with open(f) as json_file: 13 | 14 | set_ = f.split('_')[0][:-1] 15 | num = f.split('_')[0][-1] 16 | print('set: ', set_, ' num: ', num) 17 | out = [] 18 | # out_path = path_cloud +'/ply_data_'+set_+num+ '.json' 19 | out_path = path_cloud +'/'+set_+num+ '.json' 20 | print('name: ' + out_path) 21 | 22 | data = json.load(json_file) 23 | # print(data[0]) # bookshelf/bookshelf_0659.ply 24 | for path in data: 25 | clss = path.split('/')[0] 26 | name = path.split('/')[1][:-4] 27 | 28 | images = sorted(glob.glob(path_image+'/'+set_+'/model_'+name+'_*.jpg')) 29 | print(path_image+'/'+set_+'/model_'+name+'_*.jpg') 30 | images = ['/'.join(p.split('/')[-2:]) for p in images] 31 | print(len(images)) 32 | out.append(images) 33 | 34 | with open(out_path, 'w') as outfile: 35 | json.dump(out, outfile) -------------------------------------------------------------------------------- /tools/render_balls_so.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | using namespace std; 6 | 7 | struct PointInfo{ 8 | int x,y,z; 9 | float r,g,b; 10 | }; 11 | 12 | extern "C"{ 13 | 14 | void render_ball(int h,int w,unsigned char * show,int n,int * xyzs,float * c0,float * c1,float * c2,int r){ 15 | r=max(r,1); 16 | vector depth(h*w,-2100000000); 17 | vector pattern; 18 | for (int dx=-r;dx<=r;dx++) 19 | for (int dy=-r;dy<=r;dy++) 20 | if (dx*dx+dy*dy=h || y2<0 || y2>=w) && depth[x2*w+y2] 0: 87 | show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( 88 | show[:, :, 0], 1, axis=0)) 89 | if magnifyBlue >= 2: 90 | show[:, :, 0] = np.maximum(show[:, :, 0], 91 | np.roll(show[:, :, 0], -1, axis=0)) 92 | show[:, :, 0] = np.maximum(show[:, :, 0], np.roll( 93 | show[:, :, 0], 1, axis=1)) 94 | if magnifyBlue >= 2: 95 | show[:, :, 0] = np.maximum(show[:, :, 0], 96 | np.roll(show[:, :, 0], -1, axis=1)) 97 | if showrot: 98 | cv2.putText(show, 'xangle %d' % (int(xangle / np.pi * 180)), 99 | (30, showsz - 30), 0, 0.5, (255, 0, 0)) 100 | cv2.putText(show, 'yangle %d' % (int(yangle / np.pi * 180)), 101 | (30, showsz - 50), 0, 0.5, (255, 0, 0)) 102 | cv2.putText(show, 'zoom %d%%' % (int(zoom * 100)), (30, showsz - 70), 0, 103 | 0.5, (255, 0, 0)) 104 | changed = True 105 | while True: 106 | if changed: 107 | render() 108 | changed = False 109 | cv2.imshow('show3d', show) 110 | if waittime == 0: 111 | cmd = cv2.waitKey(10) % 256 112 | else: 113 | cmd = cv2.waitKey(waittime) % 256 114 | if cmd == ord('q'): 115 | break 116 | elif cmd == ord('Q'): 117 | sys.exit(0) 118 | 119 | if cmd == ord('t') or cmd == ord('p'): 120 | if cmd == ord('t'): 121 | if c_gt is None: 122 | c0 = np.zeros((len(xyz), ), dtype='float32') + 255 123 | c1 = np.zeros((len(xyz), ), dtype='float32') + 255 124 | c2 = np.zeros((len(xyz), ), dtype='float32') + 255 125 | else: 126 | c0 = c_gt[:, 0] 127 | c1 = c_gt[:, 1] 128 | c2 = c_gt[:, 2] 129 | else: 130 | if c_pred is None: 131 | c0 = np.zeros((len(xyz), ), dtype='float32') + 255 132 | c1 = np.zeros((len(xyz), ), dtype='float32') + 255 133 | c2 = np.zeros((len(xyz), ), dtype='float32') + 255 134 | else: 135 | c0 = c_pred[:, 0] 136 | c1 = c_pred[:, 1] 137 | c2 = c_pred[:, 2] 138 | if normalizecolor: 139 | c0 /= (c0.max() + 1e-14) / 255.0 140 | c1 /= (c1.max() + 1e-14) / 255.0 141 | c2 /= (c2.max() + 1e-14) / 255.0 142 | c0 = np.require(c0, 'float32', 'C') 143 | c1 = np.require(c1, 'float32', 'C') 144 | c2 = np.require(c2, 'float32', 'C') 145 | changed = True 146 | 147 | if cmd==ord('n'): 148 | zoom*=1.1 149 | changed=True 150 | elif cmd==ord('m'): 151 | zoom/=1.1 152 | changed=True 153 | elif cmd==ord('r'): 154 | zoom=1.0 155 | changed=True 156 | elif cmd==ord('s'): 157 | cv2.imwrite('show3d.png',show) 158 | if waittime!=0: 159 | break 160 | return cmd -------------------------------------------------------------------------------- /tsne.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy import array 3 | from matplotlib import pyplot as plt 4 | from sklearn.manifold import TSNE 5 | from matplotlib import cm 6 | from matplotlib.colors import ListedColormap, LinearSegmentedColormap 7 | # from tsne import bh_sne 8 | 9 | # X = np.load('./extracted_features/test_sv180_img_feat.npy') 10 | X_img = np.load('./extracted_features/ModelNet40-test-2_cm_sv180_img_feat.npy') 11 | X_pt = np.load('./extracted_features/ModelNet40-test-2_cm_sv180_cloud_feat.npy') 12 | 13 | y = np.load('./extracted_features/ModelNet40-test-2_cm_sv180_label.npy') 14 | 15 | 16 | # tsne = TSNE(n_components=2, random_state=0) 17 | tsne = TSNE(n_components=2, perplexity=15, learning_rate=10) 18 | 19 | X_img_2d = tsne.fit_transform(X_img) 20 | X_pt_2d = tsne.fit_transform(X_pt) 21 | 22 | 23 | target_ids = range(len(y)) 24 | 25 | # plt.figure(figsize=(6, 12)) 26 | 27 | fig, (ax1, ax2) = plt.subplots(1,2) 28 | 29 | fig.suptitle('Image Features and Point Cloud Features') 30 | 31 | cmap = plt.get_cmap('gnuplot') 32 | colors = [cmap(i) for i in np.linspace(0, 1, 40)] 33 | print(len(colors)) 34 | print(len(colors)) 35 | print(len(colors)) 36 | print(len(colors)) 37 | print(len(colors)) 38 | print(len(colors)) 39 | print(len(colors)) 40 | colors = ['red', 'blue', 'navy', 'green', 'violet', 'brown', 'gold', 'lime', 'teal', 'olive'] 41 | # c=np.random.rand(1,3) 42 | for i in target_ids: 43 | if i<10: 44 | ax1.scatter(X_img_2d[y == i, 0], X_img_2d[y == i, 1], c=colors[i]) 45 | ax2.scatter(X_pt_2d[y == i, 0], X_pt_2d[y == i, 1], c=colors[i]) 46 | plt.show() 47 | -------------------------------------------------------------------------------- /wandb/settings: -------------------------------------------------------------------------------- 1 | [default] 2 | entity = zhiminc 3 | project = multisemi 4 | base_url = https://api.wandb.ai 5 | 6 | --------------------------------------------------------------------------------