3D MinkowskiEngine稀疏模式重建

本文看一个简单的演示示例，该示例训练一个3D卷积神经网络，该网络用一个热点向量one-hot vector重构3D稀疏模式。这类似于Octree生成网络ICCV'17。输入的one-hot vector一热向量，来自ModelNet40数据集的3D计算机辅助设计（CAD）椅子索引。

使用MinkowskiEngine.MinkowskiConvolutionTranspose和 MinkowskiEngine.MinkowskiPruning，依次将体素上采样2倍，然后删除一些上采样的体素，以生成目标形状。常规的网络体系结构看起来类似于下图，但是细节可能有所不同。

在继续之前，请先阅读训练和数据加载。

创建稀疏模式重建网络

要从矢量创建3D网格世界中定义的稀疏张量，需要从 1×1×1分辨率体素。本文使用一个由块MinkowskiEngine.MinkowskiConvolutionTranspose，MinkowskiEngine.MinkowskiConvolution和MinkowskiEngine.MinkowskiPruning。

在前进过程forward pass中，为1）主要特征和2）稀疏体素分类创建两条路径，以删除不必要的体素。

out = upsample_block(z)

out_cls = classification(out).F

out = pruning(out, out_cls > 0)

在输入的稀疏张量达到目标分辨率之前，网络会重复执行一系列的上采样和修剪操作，以去除不必要的体素。在下图上可视化结果。注意，最终的重建非常精确地捕获了目标几何体。还可视化了上采样和修剪的分层重建过程。

运行示例

要训​​练网络，请转到Minkowski Engine根目录，然后键入：

python -m examples.reconstruction --train

要可视化网络预测或尝试预先训练的模型，请输入：

python -m examples.reconstruction

该程序将可视化两个3D形状。左边的一个是目标3D形状，右边的一个是重构的网络预测。

完整的代码可以在example / reconstruction.py找到。

import os

 

import sys

 

import subprocess

 

import argparse

 

import logging

 

import glob

 

import numpy as np

 

from time import time

 

import urllib

 

# Must be imported before large libs

 

try:

 

import open3d as o3d

 

except ImportError:

 

raise ImportError('Please install open3d and scipy with `pip install open3d scipy`.')

 

import torch

 

import torch.nn as nn

 

import torch.utils.data

 

import torch.optim as optim

 

import MinkowskiEngine as ME

 

from examples.modelnet40 import InfSampler, resample_mesh

 

M = np.array([[0.80656762, -0.5868724, -0.07091862],

 

[0.3770505, 0.418344, 0.82632997],

 

[-0.45528188, -0.6932309, 0.55870326]])

 

assert int(

 

o3d.__version__.split('.')[1]

 

) >= 8, f'Requires open3d version >= 0.8, the current version is {o3d.__version__}'

 

if not os.path.exists('ModelNet40'):

 

logging.info('Downloading the fixed ModelNet40 dataset…')

 

subprocess.run(["sh", "./examples/download_modelnet40.sh"])

 

###############################################################################

 

# Utility functions

 

###############################################################################

 

def PointCloud(points, colors=None):

 

pcd = o3d.geometry.PointCloud()

 

pcd.points = o3d.utility.Vector3dVector(points)

 

if colors is not None:

 

pcd.colors = o3d.utility.Vector3dVector(colors)

 

return pcd

 

def collate_pointcloud_fn(list_data):

 

coords, feats, labels = list(zip(*list_data))

 

# Concatenate all lists

 

return {

 

'coords': coords,

 

'xyzs': [torch.from_numpy(feat).float() for feat in feats],

 

'labels': torch.LongTensor(labels),

 

}

 

class ModelNet40Dataset(torch.utils.data.Dataset):

 

def __init__(self, phase, transform=None, config=None):

 

self.phase = phase

 

self.files = []

 

self.cache = {}

 

self.data_objects = []

 

self.transform = transform

 

self.resolution = config.resolution

 

self.last_cache_percent = 0

 

self.root = './ModelNet40'

 

fnames = glob.glob(os.path.join(self.root, 'chair/train/*.off'))

 

fnames = sorted([os.path.relpath(fname, self.root) for fname in fnames])

 

self.files = fnames

 

assert len(self.files) > 0, "No file loaded"

 

logging.info(

 

f"Loading the subset {phase} from {self.root} with {len(self.files)} files"

 

)

 

self.density = 30000

 

# Ignore warnings in obj loader

 

o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Error)

 

def __len__(self):

 

return len(self.files)

 

def __getitem__(self, idx):

 

mesh_file = os.path.join(self.root, self.files[idx])

 

if idx in self.cache:

 

xyz = self.cache[idx]

 

else:

 

# Load a mesh, over sample, copy, rotate, voxelization

 

assert os.path.exists(mesh_file)

 

pcd = o3d.io.read_triangle_mesh(mesh_file)

 

# Normalize to fit the mesh inside a unit cube while preserving aspect ratio

 

vertices = np.asarray(pcd.vertices)

 

vmax = vertices.max(0, keepdims=True)

 

vmin = vertices.min(0, keepdims=True)

 

pcd.vertices = o3d.utility.Vector3dVector(

 

(vertices - vmin) / (vmax - vmin).max())

 

# Oversample points and copy

 

xyz = resample_mesh(pcd, density=self.density)

 

self.cache[idx] = xyz

 

cache_percent = int((len(self.cache) / len(self)) * 100)

 

if cache_percent > 0 and cache_percent % 10 == 0 and cache_percent != self.last_cache_percent:

 

logging.info(

 

f"Cached {self.phase}: {len(self.cache)} / {len(self)}: {cache_percent}%"

 

)

 

self.last_cache_percent = cache_percent

 

# Use color or other features if available

 

feats = np.ones((len(xyz), 1))

 

if len(xyz) < 1000:

 

logging.info(

 

f"Skipping {mesh_file}: does not have sufficient CAD sampling density after resampling: {len(xyz)}."

 

)

 

return None

 

if self.transform:

 

xyz, feats = self.transform(xyz, feats)

 

# Get coords

 

xyz = xyz * self.resolution

 

coords = np.floor(xyz)

 

inds = ME.utils.sparse_quantize(coords, return_index=True)

 

return (coords[inds], xyz[inds], idx)

 

def make_data_loader(phase, augment_data, batch_size, shuffle, num_workers,

 

repeat, config):

 

dset = ModelNet40Dataset(phase, config=config)

 

args = {

 

'batch_size': batch_size,

 

'num_workers': num_workers,

 

'collate_fn': collate_pointcloud_fn,

 

'pin_memory': False,

 

'drop_last': False

 

}

 

if repeat:

 

args['sampler'] = InfSampler(dset, shuffle)

 

else:

 

args['shuffle'] = shuffle

 

loader = torch.utils.data.DataLoader(dset, **args)

 

return loader

 

ch = logging.StreamHandler(sys.stdout)

 

logging.getLogger().setLevel(logging.INFO)

 

logging.basicConfig(

 

format=os.uname()[1].split('.')[0] + ' %(asctime)s %(message)s',

 

datefmt='%m/%d %H:%M:%S',

 

handlers=[ch])

 

parser = argparse.ArgumentParser()

 

parser.add_argument('--resolution', type=int, default=128)

 

parser.add_argument('--max_iter', type=int, default=30000)

 

parser.add_argument('--val_freq', type=int, default=1000)

 

parser.add_argument('--batch_size', default=16, type=int)

 

parser.add_argument('--lr', default=1e-2, type=float)

 

parser.add_argument('--momentum', type=float, default=0.9)

 

parser.add_argument('--weight_decay', type=float, default=1e-4)

 

parser.add_argument('--num_workers', type=int, default=1)

 

parser.add_argument('--stat_freq', type=int, default=50)

 

parser.add_argument(

 

'--weights', type=str, default='modelnet_reconstruction.pth')

 

parser.add_argument('--load_optimizer', type=str, default='true')

 

parser.add_argument('--train', action='store_true')

 

parser.add_argument('--max_visualization', type=int, default=4)

 

###############################################################################

 

# End of utility functions

 

###############################################################################

 

class GenerativeNet(nn.Module):

 

CHANNELS = [1024, 512, 256, 128, 64, 32, 16]

 

def __init__(self, resolution, in_nchannel=512):

 

nn.Module.__init__(self)

 

self.resolution = resolution

 

# Input sparse tensor must have tensor stride 128.

 

ch = self.CHANNELS

 

# Block 1

 

self.block1 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

in_nchannel,

 

ch[0],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[0]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[0], ch[0], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[0]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolutionTranspose(

 

ch[0],

 

ch[1],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[1]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[1], ch[1], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[1]),

 

ME.MinkowskiELU(),

 

)

 

self.block1_cls = ME.MinkowskiConvolution(

 

ch[1], 1, kernel_size=1, has_bias=True, dimension=3)

 

# Block 2

 

self.block2 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

ch[1],

 

ch[2],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[2]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[2], ch[2], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[2]),

 

ME.MinkowskiELU(),

 

)

 

self.block2_cls = ME.MinkowskiConvolution(

 

ch[2], 1, kernel_size=1, has_bias=True, dimension=3)

 

# Block 3

 

self.block3 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

ch[2],

 

ch[3],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[3]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[3], ch[3], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[3]),

 

ME.MinkowskiELU(),

 

)

 

self.block3_cls = ME.MinkowskiConvolution(

 

ch[3], 1, kernel_size=1, has_bias=True, dimension=3)

 

# Block 4

 

self.block4 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

ch[3],

 

ch[4],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[4]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[4], ch[4], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[4]),

 

ME.MinkowskiELU(),

 

)

 

self.block4_cls = ME.MinkowskiConvolution(

 

ch[4], 1, kernel_size=1, has_bias=True, dimension=3)

 

# Block 5

 

self.block5 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

ch[4],

 

ch[5],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[5]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[5], ch[5], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[5]),

 

ME.MinkowskiELU(),

 

)

 

self.block5_cls = ME.MinkowskiConvolution(

 

ch[5], 1, kernel_size=1, has_bias=True, dimension=3)

 

# Block 6

 

self.block6 = nn.Sequential(

 

ME.MinkowskiConvolutionTranspose(

 

ch[5],

 

ch[6],

 

kernel_size=2,

 

stride=2,

 

generate_new_coords=True,

 

dimension=3),

 

ME.MinkowskiBatchNorm(ch[6]),

 

ME.MinkowskiELU(),

 

ME.MinkowskiConvolution(ch[6], ch[6], kernel_size=3, dimension=3),

 

ME.MinkowskiBatchNorm(ch[6]),

 

ME.MinkowskiELU(),

 

)

 

self.block6_cls = ME.MinkowskiConvolution(

 

ch[6], 1, kernel_size=1, has_bias=True, dimension=3)

 

# pruning

 

self.pruning = ME.MinkowskiPruning()

 

def get_batch_indices(self, out):

 

return out.coords_man.get_row_indices_per_batch(out.coords_key)

 

def get_target(self, out, target_key, kernel_size=1):

 

with torch.no_grad():

 

target = torch.zeros(len(out), dtype=torch.bool)

 

cm = out.coords_man

 

strided_target_key = cm.stride(

 

target_key, out.tensor_stride[0], force_creation=True)

 

ins, outs = cm.get_kernel_map(

 

out.coords_key,

 

strided_target_key,

 

kernel_size=kernel_size,

 

region_type=1)

 

for curr_in in ins:

 

target[curr_in] = 1

 

return target

 

def valid_batch_map(self, batch_map):

 

for b in batch_map:

 

if len(b) == 0:

 

return False

 

return True

 

def forward(self, z, target_key):

 

out_cls, targets = [], []

 

# Block1

 

out1 = self.block1(z)

 

out1_cls = self.block1_cls(out1)

 

target = self.get_target(out1, target_key)

 

targets.append(target)

 

out_cls.append(out1_cls)

 

keep1 = (out1_cls.F > 0).cpu().squeeze()

 

# If training, force target shape generation, use net.eval() to disable

 

if self.training:

 

keep1 += target

 

# Remove voxels 32

 

out1 = self.pruning(out1, keep1.cpu())

 

# Block 2

 

out2 = self.block2(out1)

 

out2_cls = self.block2_cls(out2)

 

target = self.get_target(out2, target_key)

 

targets.append(target)

 

out_cls.append(out2_cls)

 

keep2 = (out2_cls.F > 0).cpu().squeeze()

 

if self.training:

 

keep2 += target

 

# Remove voxels 16

 

out2 = self.pruning(out2, keep2.cpu())

 

# Block 3

 

out3 = self.block3(out2)

 

out3_cls = self.block3_cls(out3)

 

target = self.get_target(out3, target_key)

 

targets.append(target)

 

out_cls.append(out3_cls)

 

keep3 = (out3_cls.F > 0).cpu().squeeze()

 

if self.training:

 

keep3 += target

 

# Remove voxels 8

 

out3 = self.pruning(out3, keep3.cpu())

 

# Block 4

 

out4 = self.block4(out3)

 

out4_cls = self.block4_cls(out4)

 

target = self.get_target(out4, target_key)

 

targets.append(target)

 

out_cls.append(out4_cls)

 

keep4 = (out4_cls.F > 0).cpu().squeeze()

 

if self.training:

 

keep4 += target

 

# Remove voxels 4

 

out4 = self.pruning(out4, keep4.cpu())

 

# Block 5

 

out5 = self.block5(out4)

 

out5_cls = self.block5_cls(out5)

 

target = self.get_target(out5, target_key)

 

targets.append(target)

 

out_cls.append(out5_cls)

 

keep5 = (out5_cls.F > 0).cpu().squeeze()

 

if self.training:

 

keep5 += target

 

# Remove voxels 2

 

out5 = self.pruning(out5, keep5.cpu())

 

# Block 5

 

out6 = self.block6(out5)

 

out6_cls = self.block6_cls(out6)

 

target = self.get_target(out6, target_key)

 

targets.append(target)

 

out_cls.append(out6_cls)

 

keep6 = (out6_cls.F > 0).cpu().squeeze()

 

# Last layer does not require keep

 

# if self.training:

 

# keep6 += target

 

# Remove voxels 1

 

out6 = self.pruning(out6, keep6.cpu())

 

return out_cls, targets, out6

 

def train(net, dataloader, device, config):

 

in_nchannel = len(dataloader.dataset)

 

optimizer = optim.SGD(

 

net.parameters(),

 

lr=config.lr,

 

momentum=config.momentum,

 

weight_decay=config.weight_decay)

 

scheduler = optim.lr_scheduler.ExponentialLR(optimizer, 0.95)

 

crit = nn.BCEWithLogitsLoss()

 

net.train()

 

train_iter = iter(dataloader)

 

# val_iter = iter(val_dataloader)

 

logging.info(f'LR: {scheduler.get_lr()}')

 

for i in range(config.max_iter):

 

s = time()

 

data_dict = train_iter.next()

 

d = time() - s

 

optimizer.zero_grad()

 

init_coords = torch.zeros((config.batch_size, 4), dtype=torch.int)

 

init_coords[:, 0] = torch.arange(config.batch_size)

 

in_feat = torch.zeros((config.batch_size, in_nchannel))

 

in_feat[torch.arange(config.batch_size), data_dict['labels']] = 1

 

sin = ME.SparseTensor(

 

feats=in_feat,

 

coords=init_coords,

 

allow_duplicate_coords=True, # for classification, it doesn't matter

 

tensor_stride=config.resolution,

 

).to(device)

 

# Generate target sparse tensor

 

cm = sin.coords_man

 

target_key = cm.create_coords_key(

 

ME.utils.batched_coordinates(data_dict['xyzs']),

 

force_creation=True,

 

allow_duplicate_coords=True)

 

# Generate from a dense tensor

 

out_cls, targets, sout = net(sin, target_key)

 

num_layers, loss = len(out_cls), 0

 

losses = []

 

for out_cl, target in zip(out_cls, targets):

 

curr_loss = crit(out_cl.F.squeeze(),

 

target.type(out_cl.F.dtype).to(device))

 

losses.append(curr_loss.item())

 

loss += curr_loss / num_layers

 

loss.backward()

 

optimizer.step()

 

t = time() - s

 

if i % config.stat_freq == 0:

 

logging.info(

 

f'Iter: {i}, Loss: {loss.item():.3e}, Depths: {len(out_cls)} Data Loading Time: {d:.3e}, Tot Time: {t:.3e}'

 

)

 

if i % config.val_freq == 0 and i > 0:

 

torch.save(

 

{

 

'state_dict': net.state_dict(),

 

'optimizer': optimizer.state_dict(),

 

'scheduler': scheduler.state_dict(),

 

'curr_iter': i,

 

}, config.weights)

 

scheduler.step()

 

logging.info(f'LR: {scheduler.get_lr()}')

 

net.train()

 

def visualize(net, dataloader, device, config):

 

in_nchannel = len(dataloader.dataset)

 

net.eval()

 

crit = nn.BCEWithLogitsLoss()

 

n_vis = 0

 

for data_dict in dataloader:

 

init_coords = torch.zeros((config.batch_size, 4), dtype=torch.int)

 

init_coords[:, 0] = torch.arange(config.batch_size)

 

in_feat = torch.zeros((config.batch_size, in_nchannel))

 

in_feat[torch.arange(config.batch_size), data_dict['labels']] = 1

 

sin = ME.SparseTensor(

 

feats=in_feat,

 

coords=init_coords,

 

allow_duplicate_coords=True, # for classification, it doesn't matter

 

tensor_stride=config.resolution,

 

).to(device)

 

# Generate target sparse tensor

 

cm = sin.coords_man

 

target_key = cm.create_coords_key(

 

ME.utils.batched_coordinates(data_dict['xyzs']),

 

force_creation=True,

 

allow_duplicate_coords=True)

 

# Generate from a dense tensor

 

out_cls, targets, sout = net(sin, target_key)

 

num_layers, loss = len(out_cls), 0

 

for out_cl, target in zip(out_cls, targets):

 

loss += crit(out_cl.F.squeeze(),

 

target.type(out_cl.F.dtype).to(device)) / num_layers

 

batch_coords, batch_feats = sout.decomposed_coordinates_and_features

 

for b, (coords, feats) in enumerate(zip(batch_coords, batch_feats)):

 

pcd = PointCloud(coords)

 

pcd.estimate_normals()

 

pcd.translate([0.6 * config.resolution, 0, 0])

 

pcd.rotate(M)

 

opcd = PointCloud(data_dict['xyzs'][b])

 

opcd.translate([-0.6 * config.resolution, 0, 0])

 

opcd.estimate_normals()

 

opcd.rotate(M)

 

o3d.visualization.draw_geometries([pcd, opcd])

 

n_vis += 1

 

if n_vis > config.max_visualization:

 

return

 

if __name__ == '__main__':

 

config = parser.parse_args()

 

logging.info(config)

 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

 

dataloader = make_data_loader(

 

'val',

 

augment_data=True,

 

batch_size=config.batch_size,

 

shuffle=True,

 

num_workers=config.num_workers,

 

repeat=True,

 

config=config)

 

in_nchannel = len(dataloader.dataset)

 

net = GenerativeNet(config.resolution, in_nchannel=in_nchannel)

 

net.to(device)

 

logging.info(net)

 

if config.train:

 

train(net, dataloader, device, config)

 

else:

 

if not os.path.exists(config.weights):

 

logging.info(

 

f'Downloaing pretrained weights. This might take a while…')

 

urllib.request.urlretrieve(

 

"https://bit.ly/36d9m1n", filename=config.weights)

 

logging.info(f'Loading weights from {config.weights}')

 

checkpoint = torch.load(config.weights)

 

net.load_state_dict(checkpoint['state_dict'])

 

visualize(net, dataloader, device, config)