├── .gitignore ├── LICENSE ├── README.md ├── build.sh ├── datasets.py ├── download.py ├── download.sh ├── misc ├── o3d_visualize.png └── show3d.png ├── open3d_pointnet_inference.ipynb ├── open3d_visualilze.py ├── pointnet.py ├── render_balls_so.cpp ├── requirements.txt ├── show3d_balls.py ├── show_cls.py ├── show_seg.py ├── train_classification.py └── train_segmentation.py /.gitignore: -------------------------------------------------------------------------------- 1 | .ipynb_checkpoints/ 2 | data 3 | *.pyc 4 | shapenetcore_partanno_segmentation_benchmark_v0/ 5 | *.so 6 | .idea* 7 | cls/ 8 | seg/ 9 | .vscode/ 10 | *.zip 11 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Fei Xia 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Open3D-PointNet 2 | 3 | [PointNet][pointnet] implementation and visualization with [Open3D][open3d], 4 | an open-source library that supports rapid development of software that deals 5 | with 3D data. As part of the Open3D ecosystem, this repository demonstrates how 6 | Open3D can be used for ML/DL research projects. 7 | 8 | This repository is forked from 9 | [`fxia22`'s PyTorch implementation](https://github.com/fxia22/pointnet.pytorch). 10 | 11 | ![seg](misc/o3d_visualize.png) 12 | 13 | # Changelog 14 | 15 | 1. Added CPU support for non-cuda-enabled devices. 16 | 2. Used Open3D point cloud loader for loading PointNet datasets (`datasets.py`). 17 | 3. Added example for PointNet inference with Open3D Jupyter visualization 18 | (`open3d_pointnet_inference.ipynb`). 19 | 4. Added example for native OpenGL visualization with Open3D (`open3d_visualize.py`). 20 | 21 | # Setup 22 | 23 | ```bash 24 | # Install Open3D, must be v0.4.0 or above for Jupyter support 25 | pip install open3d-python 26 | 27 | # Install PyTorch 28 | # Follow: https://pytorch.org/ 29 | 30 | # Install other dependencies 31 | pip install -r requirements.txt 32 | ``` 33 | 34 | Now, launch 35 | 36 | ```bash 37 | jupyter notebook 38 | ``` 39 | 40 | and run `open3d_pointnet_inference.ipynb`. All datasets and pre-trained models 41 | shall be downloaded automatically. If you run into issues downloading files, 42 | please run `download.py` separately. 43 | 44 | [open3d]: https://github.com/IntelVCL/Open3D 45 | [pointnet]: https://arxiv.org/abs/1612.00593 46 | -------------------------------------------------------------------------------- /build.sh: -------------------------------------------------------------------------------- 1 | g++ -std=c++11 render_balls_so.cpp -o render_balls_so.so -shared -fPIC -O2 -D_GLIBCXX_USE_CXX11_ABI=0 2 | -------------------------------------------------------------------------------- /datasets.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import torch.utils.data as data 3 | from PIL import Image 4 | import os 5 | import os.path 6 | import errno 7 | import torch 8 | import json 9 | import codecs 10 | import numpy as np 11 | import progressbar 12 | import sys 13 | import torchvision.transforms as transforms 14 | import argparse 15 | import json 16 | import open3d 17 | 18 | 19 | class PartDataset(data.Dataset): 20 | def __init__(self, root, npoints = 2500, classification = False, class_choice = None, train = True): 21 | self.npoints = npoints 22 | self.root = root 23 | self.catfile = os.path.join(self.root, 'synsetoffset2category.txt') 24 | self.cat = {} 25 | 26 | self.classification = classification 27 | 28 | with open(self.catfile, 'r') as f: 29 | for line in f: 30 | ls = line.strip().split() 31 | self.cat[ls[0]] = ls[1] 32 | #print(self.cat) 33 | if not class_choice is None: 34 | self.cat = {k:v for k,v in self.cat.items() if k in class_choice} 35 | 36 | self.meta = {} 37 | for item in self.cat: 38 | #print('category', item) 39 | self.meta[item] = [] 40 | dir_point = os.path.join(self.root, self.cat[item], 'points') 41 | dir_seg = os.path.join(self.root, self.cat[item], 'points_label') 42 | #print(dir_point, dir_seg) 43 | fns = sorted(os.listdir(dir_point)) 44 | if train: 45 | fns = fns[:int(len(fns) * 0.9)] 46 | else: 47 | fns = fns[int(len(fns) * 0.9):] 48 | 49 | #print(os.path.basename(fns)) 50 | for fn in fns: 51 | token = (os.path.splitext(os.path.basename(fn))[0]) 52 | self.meta[item].append((os.path.join(dir_point, token + '.pts'), os.path.join(dir_seg, token + '.seg'))) 53 | 54 | self.datapath = [] 55 | for item in self.cat: 56 | for fn in self.meta[item]: 57 | self.datapath.append((item, fn[0], fn[1])) 58 | 59 | 60 | self.classes = dict(zip(sorted(self.cat), range(len(self.cat)))) 61 | print(self.classes) 62 | self.num_seg_classes = 0 63 | if not self.classification: 64 | for i in range(len(self.datapath)//50): 65 | l = len(np.unique(np.loadtxt(self.datapath[i][-1]).astype(np.uint8))) 66 | if l > self.num_seg_classes: 67 | self.num_seg_classes = l 68 | #print(self.num_seg_classes) 69 | 70 | 71 | def __getitem__(self, index): 72 | fn = self.datapath[index] 73 | cls = self.classes[self.datapath[index][0]] 74 | point_set = np.asarray( 75 | open3d.read_point_cloud(fn[1], format='xyz').points, 76 | dtype=np.float32) 77 | seg = np.loadtxt(fn[2]).astype(np.int64) 78 | #print(point_set.shape, seg.shape) 79 | 80 | choice = np.random.choice(len(seg), self.npoints, replace=True) 81 | #resample 82 | point_set = point_set[choice, :] 83 | seg = seg[choice] 84 | point_set = torch.from_numpy(point_set) 85 | seg = torch.from_numpy(seg) 86 | cls = torch.from_numpy(np.array([cls]).astype(np.int64)) 87 | if self.classification: 88 | return point_set, cls 89 | else: 90 | return point_set, seg 91 | 92 | def __len__(self): 93 | return len(self.datapath) 94 | 95 | 96 | if __name__ == '__main__': 97 | print('test') 98 | d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', class_choice = ['Chair']) 99 | print(len(d)) 100 | ps, seg = d[0] 101 | print(ps.size(), ps.type(), seg.size(),seg.type()) 102 | 103 | d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True) 104 | print(len(d)) 105 | ps, cls = d[0] 106 | print(ps.size(), ps.type(), cls.size(),cls.type()) 107 | -------------------------------------------------------------------------------- /download.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | # Copyright (c) 2017 Computer Vision Center (CVC) at the Universitat Autonoma de 4 | # Barcelona (UAB). 5 | # 6 | # This work is licensed under the terms of the MIT license. 7 | # For a copy, see . 8 | 9 | """Download big files from Google Drive.""" 10 | 11 | import shutil 12 | import sys 13 | import requests 14 | import os 15 | import time 16 | import urllib.request 17 | import zipfile 18 | 19 | def reporthook(count, block_size, total_size): 20 | global start_time 21 | if count == 0: 22 | start_time = time.time() 23 | return 24 | duration = time.time() - start_time 25 | progress_size = int(count * block_size) 26 | speed = int(progress_size / (1024 * duration)) 27 | percent = int(count * block_size * 100 / total_size) 28 | 29 | if percent % 5 == 0: 30 | sys.stdout.write("\r...%d%%, %d MB, %d KB/s, %d seconds passed" % 31 | (percent, progress_size / (1024 * 1024), speed, duration)) 32 | sys.stdout.flush() 33 | 34 | def sizeof_fmt(num, suffix='B'): 35 | # https://stackoverflow.com/a/1094933/5308925 36 | for unit in ['','K','M','G','T','P','E','Z']: 37 | if abs(num) < 1000.0: 38 | return "%3.2f%s%s" % (num, unit, suffix) 39 | num /= 1000.0 40 | return "%.2f%s%s" % (num, 'Yi', suffix) 41 | 42 | 43 | def print_status(destination, progress): 44 | message = "Downloading %s... %s" % (destination, sizeof_fmt(progress)) 45 | empty_space = shutil.get_terminal_size((80, 20)).columns - len(message) 46 | sys.stdout.write('\r' + message + empty_space * ' ') 47 | sys.stdout.flush() 48 | 49 | 50 | def download_file_from_google_drive(id, destination): 51 | # https://stackoverflow.com/a/39225039/5308925 52 | 53 | def save_response_content(response, destination): 54 | chunk_size = 32768 55 | written_size = 0 56 | 57 | with open(destination, "wb") as f: 58 | for chunk in response.iter_content(chunk_size): 59 | if chunk: # filter out keep-alive new chunks 60 | f.write(chunk) 61 | written_size += chunk_size 62 | print_status(destination, written_size) 63 | print('Done.') 64 | 65 | def get_confirm_token(response): 66 | for key, value in response.cookies.items(): 67 | if key.startswith('download_warning'): 68 | return value 69 | 70 | return None 71 | 72 | url = "https://docs.google.com/uc?export=download" 73 | 74 | session = requests.Session() 75 | 76 | response = session.get(url, params={'id': id}, stream=True) 77 | token = get_confirm_token(response) 78 | 79 | if token: 80 | params = {'id': id, 'confirm': token} 81 | response = session.get(url, params=params, stream=True) 82 | 83 | save_response_content(response, destination) 84 | 85 | 86 | def download_contents(): 87 | # download model 88 | model_path = './cls_model.pth' 89 | if os.path.isfile(model_path): 90 | print('Model file already downloaded in', model_path) 91 | else: 92 | download_file_from_google_drive('1WWf5B5fmik5_P1dwxltJ-atRkYeCcCC5', './cls_model.pth') 93 | 94 | # download dataset 95 | dataset_path = './shapenetcore_partanno_segmentation_benchmark_v0.zip' 96 | if os.path.isfile(dataset_path): 97 | print('Dataset file already downloaded in', dataset_path) 98 | else: 99 | dataset_url = 'https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_segmentation_benchmark_v0.zip' 100 | urllib.request.urlretrieve(dataset_url, os.path.basename(dataset_url), reporthook) 101 | 102 | # unzip dataset 103 | zip_ref = zipfile.ZipFile(os.path.basename(dataset_url), 'r') 104 | zip_ref.extractall('.') 105 | zip_ref.close() 106 | 107 | print('Now unzipping...Wait for 2 minutes ish...!') 108 | return 0 109 | 110 | if __name__ == '__main__': 111 | download_contents() 112 | -------------------------------------------------------------------------------- /download.sh: -------------------------------------------------------------------------------- 1 | wget https://shapenet.cs.stanford.edu/ericyi/shapenetcore_partanno_segmentation_benchmark_v0.zip --no-check-certificate 2 | unzip shapenetcore_partanno_segmentation_benchmark_v0.zip 3 | rm shapenetcore_partanno_segmentation_benchmark_v0.zip 4 | -------------------------------------------------------------------------------- /misc/o3d_visualize.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/isl-org/Open3D-PointNet/d5c0d146f1f529a43704923c8895ba4b36d1781b/misc/o3d_visualize.png -------------------------------------------------------------------------------- /misc/show3d.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/isl-org/Open3D-PointNet/d5c0d146f1f529a43704923c8895ba4b36d1781b/misc/show3d.png -------------------------------------------------------------------------------- /open3d_pointnet_inference.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import os\n", 10 | "import random\n", 11 | "from random import randrange\n", 12 | "from IPython.display import clear_output\n", 13 | "import numpy as np\n", 14 | "import torch\n", 15 | "import torch.nn as nn\n", 16 | "import torch.nn.parallel\n", 17 | "import torch.optim as optim\n", 18 | "import torch.utils.data\n", 19 | "import torchvision.datasets as dset\n", 20 | "import torchvision.transforms as transforms\n", 21 | "import torchvision.utils as vutils\n", 22 | "from torch.autograd import Variable\n", 23 | "from datasets import PartDataset\n", 24 | "from pointnet import PointNetCls\n", 25 | "import torch.nn.functional as F\n", 26 | "import matplotlib.pyplot as plt\n", 27 | "import open3d as o3\n", 28 | "import download\n", 29 | "from open3d import JVisualizer\n", 30 | "if torch.cuda.is_available():\n", 31 | " import torch.backends.cudnn as cudnn" 32 | ] 33 | }, 34 | { 35 | "cell_type": "code", 36 | "execution_count": null, 37 | "metadata": {}, 38 | "outputs": [], 39 | "source": [ 40 | "# General parameters\n", 41 | "NUM_POINTS = 10000\n", 42 | "MODEL_PATH = './cls_model.pth'\n", 43 | "DATA_FOLDER = './shapenetcore_partanno_segmentation_benchmark_v0'\n", 44 | "\n", 45 | "# download dataset and pre-trained model\n", 46 | "download.download_contents()" 47 | ] 48 | }, 49 | { 50 | "cell_type": "code", 51 | "execution_count": null, 52 | "metadata": {}, 53 | "outputs": [], 54 | "source": [ 55 | "# Create dataset object\n", 56 | "test_dataset_seg = PartDataset(\n", 57 | " root=DATA_FOLDER,\n", 58 | " train=False,\n", 59 | " classification=False,\n", 60 | " npoints=NUM_POINTS)\n", 61 | "\n", 62 | "# Problem ontology\n", 63 | "classes_dict = {'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4, \n", 64 | " 'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9,\n", 65 | " 'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13, \n", 66 | " 'Skateboard': 14, 'Table': 15}" 67 | ] 68 | }, 69 | { 70 | "cell_type": "code", 71 | "execution_count": null, 72 | "metadata": {}, 73 | "outputs": [], 74 | "source": [ 75 | "# Create the classification network from pre-trained model\n", 76 | "classifier = PointNetCls(k=len(classes_dict.items()), num_points=NUM_POINTS)\n", 77 | "if torch.cuda.is_available():\n", 78 | " classifier.cuda()\n", 79 | " classifier.load_state_dict(torch.load(MODEL_PATH))\n", 80 | "else:\n", 81 | " classifier.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))\n", 82 | "classifier.eval()" 83 | ] 84 | }, 85 | { 86 | "cell_type": "code", 87 | "execution_count": null, 88 | "metadata": {}, 89 | "outputs": [], 90 | "source": [ 91 | "# Simple point cloud coloring mapping\n", 92 | "def read_pointnet_colors(seg_labels):\n", 93 | " map_label_to_rgb = {\n", 94 | " 1: [0, 255, 0],\n", 95 | " 2: [0, 0, 255],\n", 96 | " 3: [255, 0, 0],\n", 97 | " 4: [255, 0, 255], # purple\n", 98 | " 5: [0, 255, 255], # cyan\n", 99 | " 6: [255, 255, 0], # yellow\n", 100 | " }\n", 101 | " colors = np.array([map_label_to_rgb[label] for label in seg_labels])\n", 102 | " return colors" 103 | ] 104 | }, 105 | { 106 | "cell_type": "code", 107 | "execution_count": null, 108 | "metadata": {}, 109 | "outputs": [], 110 | "source": [ 111 | "# Three.js based visualizer\n", 112 | "visualizer = JVisualizer()\n", 113 | "\n", 114 | "# Basic inference and visualization loop\n", 115 | "MAX_SAMPLES = 15\n", 116 | "for samples in range(MAX_SAMPLES):\n", 117 | " random_index = randrange(len(test_dataset_seg))\n", 118 | " print('[Sample {} / {}]'.format(random_index, len(test_dataset_seg)))\n", 119 | " \n", 120 | " # clean visualization\n", 121 | " visualizer.clear()\n", 122 | " clear_output()\n", 123 | "\n", 124 | " # get next sample\n", 125 | " point_set, seg = test_dataset_seg.__getitem__(random_index)\n", 126 | "\n", 127 | " # create cloud for visualization\n", 128 | " cloud = o3.PointCloud()\n", 129 | " cloud.points = o3.Vector3dVector(point_set)\n", 130 | " cloud.colors = o3.Vector3dVector(read_pointnet_colors(seg.numpy()))\n", 131 | "\n", 132 | " # perform inference in GPU\n", 133 | " points = Variable(point_set.unsqueeze(0))\n", 134 | " points = points.transpose(2, 1)\n", 135 | " if torch.cuda.is_available():\n", 136 | " points = points.cuda()\n", 137 | " pred_logsoft, _ = classifier(points)\n", 138 | "\n", 139 | " # move data back to cpu for visualization\n", 140 | " pred_logsoft_cpu = pred_logsoft.data.cpu().numpy().squeeze()\n", 141 | " pred_soft_cpu = np.exp(pred_logsoft_cpu)\n", 142 | " pred_class = np.argmax(pred_soft_cpu)\n", 143 | "\n", 144 | " # let's visualize the input sample\n", 145 | " visualizer.add_geometry(cloud)\n", 146 | " visualizer.show()\n", 147 | " \n", 148 | " # Visualize probabilities\n", 149 | " plt.xticks(list(classes_dict.values()), list(classes_dict.keys()),rotation=90)\n", 150 | " plt.xlabel('Classes')\n", 151 | " plt.ylabel('Probabilities')\n", 152 | " plt.plot(list(classes_dict.values()), pred_soft_cpu)\n", 153 | " plt.show()\n", 154 | "\n", 155 | " input('Your object is a [{}] with probability {:0.3}. Press enter to continue!'\n", 156 | " .format(list(classes_dict.keys())[pred_class], pred_soft_cpu[pred_class]))\n", 157 | " " 158 | ] 159 | } 160 | ], 161 | "metadata": { 162 | "kernelspec": { 163 | "display_name": "Python 3", 164 | "language": "python", 165 | "name": "python3" 166 | }, 167 | "language_info": { 168 | "codemirror_mode": { 169 | "name": "ipython", 170 | "version": 3 171 | }, 172 | "file_extension": ".py", 173 | "mimetype": "text/x-python", 174 | "name": "python", 175 | "nbconvert_exporter": "python", 176 | "pygments_lexer": "ipython3", 177 | "version": "3.6.5" 178 | } 179 | }, 180 | "nbformat": 4, 181 | "nbformat_minor": 2 182 | } 183 | -------------------------------------------------------------------------------- /open3d_visualilze.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import open3d 4 | 5 | 6 | class Visualizer: 7 | dataset_path = 'shapenetcore_partanno_segmentation_benchmark_v0' 8 | map_label_to_rgb = { 9 | 1: [0, 255, 0], 10 | 2: [0, 0, 255], 11 | 3: [255, 0, 0], 12 | 4: [255, 0, 255], # purple 13 | 5: [0, 255, 255], # cyan 14 | 6: [255, 255, 0], # yellow 15 | } 16 | 17 | def __init__(self): 18 | pass 19 | 20 | def visualize(self, obj_category, obj_id): 21 | # Concat paths 22 | pts_path = os.path.join(Visualizer.dataset_path, obj_category, 23 | 'points', obj_id + '.pts') 24 | label_path = os.path.join(Visualizer.dataset_path, obj_category, 25 | 'points_label', obj_id + '.seg') 26 | 27 | # Read point cloud 28 | point_cloud = open3d.read_point_cloud(pts_path, format='xyz') 29 | print(point_cloud) 30 | 31 | # Read label and map to color 32 | labels = np.loadtxt(label_path) 33 | colors = np.array( 34 | [Visualizer.map_label_to_rgb[label] for label in labels]) 35 | point_cloud.colors = open3d.Vector3dVector(colors) 36 | open3d.draw_geometries([point_cloud]) 37 | 38 | 39 | if __name__ == '__main__': 40 | v = Visualizer() 41 | v.visualize('02691156', '1a04e3eab45ca15dd86060f189eb133') 42 | -------------------------------------------------------------------------------- /pointnet.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import os 4 | import random 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.parallel 8 | import torch.optim as optim 9 | import torch.utils.data 10 | import torchvision.transforms as transforms 11 | import torchvision.utils as vutils 12 | from torch.autograd import Variable 13 | from PIL import Image 14 | import numpy as np 15 | import matplotlib.pyplot as plt 16 | import pdb 17 | import torch.nn.functional as F 18 | if torch.cuda.is_available(): 19 | import torch.backends.cudnn as cudnn 20 | 21 | 22 | class STN3d(nn.Module): 23 | def __init__(self, num_points = 2500): 24 | super(STN3d, self).__init__() 25 | self.num_points = num_points 26 | self.conv1 = torch.nn.Conv1d(3, 64, 1) 27 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 28 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 29 | self.mp1 = torch.nn.MaxPool1d(num_points) 30 | self.fc1 = nn.Linear(1024, 512) 31 | self.fc2 = nn.Linear(512, 256) 32 | self.fc3 = nn.Linear(256, 9) 33 | self.relu = nn.ReLU() 34 | 35 | self.bn1 = nn.BatchNorm1d(64) 36 | self.bn2 = nn.BatchNorm1d(128) 37 | self.bn3 = nn.BatchNorm1d(1024) 38 | self.bn4 = nn.BatchNorm1d(512) 39 | self.bn5 = nn.BatchNorm1d(256) 40 | 41 | 42 | def forward(self, x): 43 | batchsize = x.size()[0] 44 | x = F.relu(self.bn1(self.conv1(x))) 45 | x = F.relu(self.bn2(self.conv2(x))) 46 | x = F.relu(self.bn3(self.conv3(x))) 47 | x = self.mp1(x) 48 | x = x.view(-1, 1024) 49 | 50 | x = F.relu(self.bn4(self.fc1(x))) 51 | x = F.relu(self.bn5(self.fc2(x))) 52 | x = self.fc3(x) 53 | 54 | iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1) 55 | if x.is_cuda: 56 | iden = iden.cuda() 57 | x = x + iden 58 | x = x.view(-1, 3, 3) 59 | return x 60 | 61 | 62 | class PointNetfeat(nn.Module): 63 | def __init__(self, num_points = 2500, global_feat = True): 64 | super(PointNetfeat, self).__init__() 65 | self.stn = STN3d(num_points = num_points) 66 | self.conv1 = torch.nn.Conv1d(3, 64, 1) 67 | self.conv2 = torch.nn.Conv1d(64, 128, 1) 68 | self.conv3 = torch.nn.Conv1d(128, 1024, 1) 69 | self.bn1 = nn.BatchNorm1d(64) 70 | self.bn2 = nn.BatchNorm1d(128) 71 | self.bn3 = nn.BatchNorm1d(1024) 72 | self.mp1 = torch.nn.MaxPool1d(num_points) 73 | self.num_points = num_points 74 | self.global_feat = global_feat 75 | def forward(self, x): 76 | batchsize = x.size()[0] 77 | trans = self.stn(x) 78 | x = x.transpose(2,1) 79 | x = torch.bmm(x, trans) 80 | x = x.transpose(2,1) 81 | x = F.relu(self.bn1(self.conv1(x))) 82 | pointfeat = x 83 | x = F.relu(self.bn2(self.conv2(x))) 84 | x = self.bn3(self.conv3(x)) 85 | x = self.mp1(x) 86 | x = x.view(-1, 1024) 87 | if self.global_feat: 88 | return x, trans 89 | else: 90 | x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points) 91 | return torch.cat([x, pointfeat], 1), trans 92 | 93 | class PointNetCls(nn.Module): 94 | def __init__(self, num_points = 2500, k = 2): 95 | super(PointNetCls, self).__init__() 96 | self.num_points = num_points 97 | self.feat = PointNetfeat(num_points, global_feat=True) 98 | self.fc1 = nn.Linear(1024, 512) 99 | self.fc2 = nn.Linear(512, 256) 100 | self.fc3 = nn.Linear(256, k) 101 | self.bn1 = nn.BatchNorm1d(512) 102 | self.bn2 = nn.BatchNorm1d(256) 103 | self.relu = nn.ReLU() 104 | def forward(self, x): 105 | x, trans = self.feat(x) 106 | x = F.relu(self.bn1(self.fc1(x))) 107 | x = F.relu(self.bn2(self.fc2(x))) 108 | x = self.fc3(x) 109 | return F.log_softmax(x, dim=-1), trans 110 | 111 | class PointNetDenseCls(nn.Module): 112 | def __init__(self, num_points = 2500, k = 2): 113 | super(PointNetDenseCls, self).__init__() 114 | self.num_points = num_points 115 | self.k = k 116 | self.feat = PointNetfeat(num_points, global_feat=False) 117 | self.conv1 = torch.nn.Conv1d(1088, 512, 1) 118 | self.conv2 = torch.nn.Conv1d(512, 256, 1) 119 | self.conv3 = torch.nn.Conv1d(256, 128, 1) 120 | self.conv4 = torch.nn.Conv1d(128, self.k, 1) 121 | self.bn1 = nn.BatchNorm1d(512) 122 | self.bn2 = nn.BatchNorm1d(256) 123 | self.bn3 = nn.BatchNorm1d(128) 124 | 125 | def forward(self, x): 126 | batchsize = x.size()[0] 127 | x, trans = self.feat(x) 128 | x = F.relu(self.bn1(self.conv1(x))) 129 | x = F.relu(self.bn2(self.conv2(x))) 130 | x = F.relu(self.bn3(self.conv3(x))) 131 | x = self.conv4(x) 132 | x = x.transpose(2,1).contiguous() 133 | x = F.log_softmax(x.view(-1,self.k), dim=-1) 134 | x = x.view(batchsize, self.num_points, self.k) 135 | return x, trans 136 | 137 | 138 | if __name__ == '__main__': 139 | sim_data = Variable(torch.rand(32,3,2500)) 140 | trans = STN3d() 141 | out = trans(sim_data) 142 | print('stn', out.size()) 143 | 144 | pointfeat = PointNetfeat(global_feat=True) 145 | out, _ = pointfeat(sim_data) 146 | print('global feat', out.size()) 147 | 148 | pointfeat = PointNetfeat(global_feat=False) 149 | out, _ = pointfeat(sim_data) 150 | print('point feat', out.size()) 151 | 152 | cls = PointNetCls(k = 5) 153 | out, _ = cls(sim_data) 154 | print('class', out.size()) 155 | 156 | seg = PointNetDenseCls(k = 3) 157 | out, _ = seg(sim_data) 158 | print('seg', out.size()) 159 | -------------------------------------------------------------------------------- /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: 86 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=0)) 87 | if magnifyBlue>=2: 88 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=0)) 89 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],1,axis=1)) 90 | if magnifyBlue>=2: 91 | show[:,:,0]=np.maximum(show[:,:,0],np.roll(show[:,:,0],-1,axis=1)) 92 | if showrot: 93 | cv2.putText(show,'xangle %d'%(int(xangle/np.pi*180)),(30,showsz-30),0,0.5,cv2.cv.CV_RGB(255,0,0)) 94 | cv2.putText(show,'yangle %d'%(int(yangle/np.pi*180)),(30,showsz-50),0,0.5,cv2.cv.CV_RGB(255,0,0)) 95 | cv2.putText(show,'zoom %d%%'%(int(zoom*100)),(30,showsz-70),0,0.5,cv2.cv.CV_RGB(255,0,0)) 96 | changed=True 97 | while True: 98 | if changed: 99 | render() 100 | changed=False 101 | cv2.imshow('show3d',show) 102 | if waittime==0: 103 | cmd=cv2.waitKey(10)%256 104 | else: 105 | cmd=cv2.waitKey(waittime)%256 106 | if cmd==ord('q'): 107 | break 108 | elif cmd==ord('Q'): 109 | sys.exit(0) 110 | 111 | if cmd==ord('t') or cmd == ord('p'): 112 | if cmd == ord('t'): 113 | if c_gt is None: 114 | c0=np.zeros((len(xyz),),dtype='float32')+255 115 | c1=np.zeros((len(xyz),),dtype='float32')+255 116 | c2=np.zeros((len(xyz),),dtype='float32')+255 117 | else: 118 | c0=c_gt[:,0] 119 | c1=c_gt[:,1] 120 | c2=c_gt[:,2] 121 | else: 122 | if c_pred is None: 123 | c0=np.zeros((len(xyz),),dtype='float32')+255 124 | c1=np.zeros((len(xyz),),dtype='float32')+255 125 | c2=np.zeros((len(xyz),),dtype='float32')+255 126 | else: 127 | c0=c_pred[:,0] 128 | c1=c_pred[:,1] 129 | c2=c_pred[:,2] 130 | if normalizecolor: 131 | c0/=(c0.max()+1e-14)/255.0 132 | c1/=(c1.max()+1e-14)/255.0 133 | c2/=(c2.max()+1e-14)/255.0 134 | c0=np.require(c0,'float32','C') 135 | c1=np.require(c1,'float32','C') 136 | c2=np.require(c2,'float32','C') 137 | changed = True 138 | 139 | 140 | 141 | if cmd==ord('n'): 142 | zoom*=1.1 143 | changed=True 144 | elif cmd==ord('m'): 145 | zoom/=1.1 146 | changed=True 147 | elif cmd==ord('r'): 148 | zoom=1.0 149 | changed=True 150 | elif cmd==ord('s'): 151 | cv2.imwrite('show3d.png',show) 152 | if waittime!=0: 153 | break 154 | return cmd 155 | if __name__=='__main__': 156 | np.random.seed(100) 157 | showpoints(np.random.randn(2500,3)) 158 | 159 | -------------------------------------------------------------------------------- /show_cls.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import os 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.parallel 9 | import torch.backends.cudnn as cudnn 10 | import torch.optim as optim 11 | import torch.utils.data 12 | import torchvision.datasets as dset 13 | import torchvision.transforms as transforms 14 | import torchvision.utils as vutils 15 | from torch.autograd import Variable 16 | from datasets import PartDataset 17 | from pointnet import PointNetCls 18 | import torch.nn.functional as F 19 | import matplotlib.pyplot as plt 20 | 21 | 22 | #showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500))) 23 | 24 | parser = argparse.ArgumentParser() 25 | 26 | parser.add_argument('--model', type=str, default = '', help='model path') 27 | parser.add_argument('--num_points', type=int, default=2500, help='input batch size') 28 | 29 | 30 | opt = parser.parse_args() 31 | print (opt) 32 | 33 | test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0' , train = False, classification = True, npoints = opt.num_points) 34 | 35 | testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle = True) 36 | 37 | 38 | classifier = PointNetCls(k = len(test_dataset.classes), num_points = opt.num_points) 39 | classifier.cuda() 40 | classifier.load_state_dict(torch.load(opt.model)) 41 | classifier.eval() 42 | 43 | 44 | for i, data in enumerate(testdataloader, 0): 45 | points, target = data 46 | points, target = Variable(points), Variable(target[:, 0]) 47 | points = points.transpose(2, 1) 48 | points, target = points.cuda(), target.cuda() 49 | pred, _ = classifier(points) 50 | loss = F.nll_loss(pred, target) 51 | 52 | pred_choice = pred.data.max(1)[1] 53 | correct = pred_choice.eq(target.data).cpu().sum() 54 | print('i:%d loss: %f accuracy: %f' %(i, loss.data[0], correct/float(32))) 55 | -------------------------------------------------------------------------------- /show_seg.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from show3d_balls import * 3 | import argparse 4 | import os 5 | import random 6 | import numpy as np 7 | import torch 8 | import torch.nn as nn 9 | import torch.nn.parallel 10 | import torch.backends.cudnn as cudnn 11 | import torch.optim as optim 12 | import torch.utils.data 13 | import torchvision.datasets as dset 14 | import torchvision.transforms as transforms 15 | import torchvision.utils as vutils 16 | from torch.autograd import Variable 17 | from datasets import PartDataset 18 | from pointnet import PointNetDenseCls 19 | import torch.nn.functional as F 20 | import matplotlib.pyplot as plt 21 | 22 | 23 | #showpoints(np.random.randn(2500,3), c1 = np.random.uniform(0,1,size = (2500))) 24 | 25 | parser = argparse.ArgumentParser() 26 | 27 | parser.add_argument('--model', type=str, default = '', help='model path') 28 | parser.add_argument('--idx', type=int, default = 0, help='model index') 29 | 30 | 31 | 32 | opt = parser.parse_args() 33 | print (opt) 34 | 35 | d = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', class_choice = ['Chair'], train = False) 36 | 37 | idx = opt.idx 38 | 39 | print("model %d/%d" %( idx, len(d))) 40 | 41 | point, seg = d[idx] 42 | print(point.size(), seg.size()) 43 | 44 | point_np = point.numpy() 45 | 46 | 47 | 48 | cmap = plt.cm.get_cmap("hsv", 10) 49 | cmap = np.array([cmap(i) for i in range(10)])[:,:3] 50 | gt = cmap[seg.numpy() - 1, :] 51 | 52 | classifier = PointNetDenseCls(k = 4) 53 | classifier.load_state_dict(torch.load(opt.model)) 54 | classifier.eval() 55 | 56 | point = point.transpose(1,0).contiguous() 57 | 58 | point = Variable(point.view(1, point.size()[0], point.size()[1])) 59 | pred, _ = classifier(point) 60 | pred_choice = pred.data.max(2)[1] 61 | print(pred_choice) 62 | 63 | #print(pred_choice.size()) 64 | pred_color = cmap[pred_choice.numpy()[0], :] 65 | 66 | #print(pred_color.shape) 67 | showpoints(point_np, gt, pred_color) 68 | 69 | -------------------------------------------------------------------------------- /train_classification.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import os 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.parallel 9 | import torch.optim as optim 10 | import torch.utils.data 11 | import torchvision.datasets as dset 12 | import torchvision.transforms as transforms 13 | import torchvision.utils as vutils 14 | from torch.autograd import Variable 15 | from datasets import PartDataset 16 | from pointnet import PointNetCls 17 | import torch.nn.functional as F 18 | if torch.cuda.is_available(): 19 | import torch.backends.cudnn as cudnn 20 | 21 | 22 | 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--batchSize', type=int, default=32, help='input batch size') 25 | parser.add_argument('--num_points', type=int, default=2500, help='input batch size') 26 | parser.add_argument('--workers', type=int, help='number of data loading workers', default=4) 27 | parser.add_argument('--nepoch', type=int, default=25, help='number of epochs to train for') 28 | parser.add_argument('--outf', type=str, default='cls', help='output folder') 29 | parser.add_argument('--model', type=str, default = '', help='model path') 30 | 31 | opt = parser.parse_args() 32 | print (opt) 33 | 34 | blue = lambda x:'\033[94m' + x + '\033[0m' 35 | 36 | opt.manualSeed = random.randint(1, 10000) # fix seed 37 | print("Random Seed: ", opt.manualSeed) 38 | random.seed(opt.manualSeed) 39 | torch.manual_seed(opt.manualSeed) 40 | 41 | dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True, npoints = opt.num_points) 42 | dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, 43 | shuffle=True, num_workers=int(opt.workers)) 44 | 45 | test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = True, train = False, npoints = opt.num_points) 46 | testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=opt.batchSize, 47 | shuffle=True, num_workers=int(opt.workers)) 48 | 49 | print(len(dataset), len(test_dataset)) 50 | num_classes = len(dataset.classes) 51 | print('classes', num_classes) 52 | 53 | try: 54 | os.makedirs(opt.outf) 55 | except OSError: 56 | pass 57 | 58 | 59 | classifier = PointNetCls(k = num_classes, num_points = opt.num_points) 60 | 61 | 62 | if opt.model != '': 63 | classifier.load_state_dict(torch.load(opt.model)) 64 | 65 | 66 | optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) 67 | if torch.cuda.is_available(): 68 | classifier.cuda() 69 | 70 | num_batch = len(dataset)/opt.batchSize 71 | 72 | for epoch in range(opt.nepoch): 73 | for i, data in enumerate(dataloader, 0): 74 | points, target = data 75 | points, target = Variable(points), Variable(target[:,0]) 76 | points = points.transpose(2,1) 77 | if torch.cuda.is_available(): 78 | points, target = points.cuda(), target.cuda() 79 | optimizer.zero_grad() 80 | classifier = classifier.train() 81 | pred, _ = classifier(points) 82 | loss = F.nll_loss(pred, target) 83 | loss.backward() 84 | optimizer.step() 85 | pred_choice = pred.data.max(1)[1] 86 | correct = pred_choice.eq(target.data).cpu().sum() 87 | print('[%d: %d/%d] train loss: %f accuracy: %f' %(epoch, i, num_batch, loss.item(),correct.item() / float(opt.batchSize))) 88 | 89 | if i % 10 == 0: 90 | j, data = next(enumerate(testdataloader, 0)) 91 | points, target = data 92 | points, target = Variable(points), Variable(target[:,0]) 93 | points = points.transpose(2,1) 94 | if torch.cuda.is_available(): 95 | points, target = points.cuda(), target.cuda() 96 | classifier = classifier.eval() 97 | pred, _ = classifier(points) 98 | loss = F.nll_loss(pred, target) 99 | pred_choice = pred.data.max(1)[1] 100 | correct = pred_choice.eq(target.data).cpu().sum() 101 | print('[%d: %d/%d] %s loss: %f accuracy: %f' %(epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize))) 102 | 103 | torch.save(classifier.state_dict(), '%s/cls_model_%d.pth' % (opt.outf, epoch)) 104 | -------------------------------------------------------------------------------- /train_segmentation.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import argparse 3 | import os 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.parallel 9 | import torch.optim as optim 10 | import torch.utils.data 11 | import torchvision.datasets as dset 12 | import torchvision.transforms as transforms 13 | import torchvision.utils as vutils 14 | from torch.autograd import Variable 15 | from datasets import PartDataset 16 | from pointnet import PointNetDenseCls 17 | import torch.nn.functional as F 18 | if torch.cuda.is_available(): 19 | import torch.backends.cudnn as cudnn 20 | 21 | 22 | 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--batchSize', type=int, default=32, help='input batch size') 25 | parser.add_argument('--workers', type=int, help='number of data loading workers', default=4) 26 | parser.add_argument('--nepoch', type=int, default=25, help='number of epochs to train for') 27 | parser.add_argument('--outf', type=str, default='seg', help='output folder') 28 | parser.add_argument('--model', type=str, default = '', help='model path') 29 | 30 | 31 | opt = parser.parse_args() 32 | print (opt) 33 | 34 | opt.manualSeed = random.randint(1, 10000) # fix seed 35 | print("Random Seed: ", opt.manualSeed) 36 | random.seed(opt.manualSeed) 37 | torch.manual_seed(opt.manualSeed) 38 | 39 | dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = False, class_choice = ['Chair']) 40 | dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize, 41 | shuffle=True, num_workers=int(opt.workers)) 42 | 43 | test_dataset = PartDataset(root = 'shapenetcore_partanno_segmentation_benchmark_v0', classification = False, class_choice = ['Chair'], train = False) 44 | testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=opt.batchSize, 45 | shuffle=True, num_workers=int(opt.workers)) 46 | 47 | print(len(dataset), len(test_dataset)) 48 | num_classes = dataset.num_seg_classes 49 | print('classes', num_classes) 50 | try: 51 | os.makedirs(opt.outf) 52 | except OSError: 53 | pass 54 | 55 | blue = lambda x:'\033[94m' + x + '\033[0m' 56 | 57 | 58 | classifier = PointNetDenseCls(k = num_classes) 59 | 60 | if opt.model != '': 61 | classifier.load_state_dict(torch.load(opt.model)) 62 | 63 | optimizer = optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) 64 | if torch.cuda.is_available(): 65 | classifier.cuda() 66 | 67 | num_batch = len(dataset)/opt.batchSize 68 | 69 | for epoch in range(opt.nepoch): 70 | for i, data in enumerate(dataloader, 0): 71 | points, target = data 72 | points, target = Variable(points), Variable(target) 73 | points = points.transpose(2,1) 74 | if torch.cuda.is_available(): 75 | points, target = points.cuda(), target.cuda() 76 | optimizer.zero_grad() 77 | classifier = classifier.train() 78 | pred, _ = classifier(points) 79 | pred = pred.view(-1, num_classes) 80 | target = target.view(-1,1)[:,0] - 1 81 | #print(pred.size(), target.size()) 82 | loss = F.nll_loss(pred, target) 83 | loss.backward() 84 | optimizer.step() 85 | pred_choice = pred.data.max(1)[1] 86 | correct = pred_choice.eq(target.data).cpu().sum() 87 | print('[%d: %d/%d] train loss: %f accuracy: %f' %(epoch, i, num_batch, loss.item(), correct.item()/float(opt.batchSize * 2500))) 88 | 89 | if i % 10 == 0: 90 | j, data = next(enumerate(testdataloader, 0)) 91 | points, target = data 92 | points, target = Variable(points), Variable(target) 93 | points = points.transpose(2,1) 94 | if torch.cuda.is_available(): 95 | points, target = points.cuda(), target.cuda() 96 | classifier = classifier.eval() 97 | pred, _ = classifier(points) 98 | pred = pred.view(-1, num_classes) 99 | target = target.view(-1,1)[:,0] - 1 100 | 101 | loss = F.nll_loss(pred, target) 102 | pred_choice = pred.data.max(1)[1] 103 | correct = pred_choice.eq(target.data).cpu().sum() 104 | print('[%d: %d/%d] %s loss: %f accuracy: %f' %(epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize * 2500))) 105 | 106 | torch.save(classifier.state_dict(), '%s/seg_model_%d.pth' % (opt.outf, epoch)) --------------------------------------------------------------------------------