1703 lines
63 KiB
Python
1703 lines
63 KiB
Python
#
|
|
#
|
|
# 0=================================0
|
|
# | Kernel Point Convolutions |
|
|
# 0=================================0
|
|
#
|
|
#
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Class handling NPM3D dataset.
|
|
# Implements a Dataset, a Sampler, and a collate_fn
|
|
#
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Hugues THOMAS - 11/06/2018
|
|
#
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Imports and global variables
|
|
# \**********************************/
|
|
#
|
|
|
|
# Common libs
|
|
import time
|
|
import numpy as np
|
|
import pickle
|
|
import torch
|
|
import math
|
|
import warnings
|
|
from multiprocessing import Lock
|
|
|
|
# OS functions
|
|
from os import listdir
|
|
from os.path import exists, join, isdir
|
|
|
|
# Dataset parent class
|
|
from datasets.common import PointCloudDataset
|
|
from torch.utils.data import Sampler, get_worker_info
|
|
from utils.mayavi_visu import *
|
|
|
|
from datasets.common import grid_subsampling
|
|
from utils.config import bcolors
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Dataset class definition
|
|
# \******************************/
|
|
|
|
|
|
class NPM3DDataset(PointCloudDataset):
|
|
"""Class to handle NPM3D dataset."""
|
|
|
|
def __init__(self, config, set='training', use_potentials=True, load_data=True):
|
|
"""
|
|
This dataset is small enough to be stored in-memory, so load all point clouds here
|
|
"""
|
|
PointCloudDataset.__init__(self, 'NPM3D')
|
|
|
|
############
|
|
# Parameters
|
|
############
|
|
|
|
# Dict from labels to names
|
|
self.label_to_names = {0: 'unclassified',
|
|
1: 'ground',
|
|
2: 'building',
|
|
3: 'pole', # pole - road sign - traffic light
|
|
4: 'bollard', # bollard - small pole
|
|
5: 'trash', # trash can
|
|
6: 'barrier',
|
|
7: 'pedestrian',
|
|
8: 'car',
|
|
9: 'natural' # natural - vegetation
|
|
}
|
|
|
|
# Initialize a bunch of variables concerning class labels
|
|
self.init_labels()
|
|
|
|
# List of classes ignored during training (can be empty)
|
|
self.ignored_labels = np.array([0])
|
|
|
|
# Dataset folder
|
|
self.path = '../../Data/Paris'
|
|
|
|
# Type of task conducted on this dataset
|
|
self.dataset_task = 'cloud_segmentation'
|
|
|
|
# Update number of class and data task in configuration
|
|
config.num_classes = self.num_classes - len(self.ignored_labels)
|
|
config.dataset_task = self.dataset_task
|
|
|
|
# Parameters from config
|
|
self.config = config
|
|
|
|
# Training or test set
|
|
self.set = set
|
|
|
|
# Using potential or random epoch generation
|
|
self.use_potentials = use_potentials
|
|
|
|
# Path of the training files
|
|
# self.train_path = 'original_ply'
|
|
self.train_path = 'train'
|
|
self.original_ply_path = 'original_ply'
|
|
|
|
# List of files to process
|
|
ply_path = join(self.path, self.train_path)
|
|
|
|
# Proportion of validation scenes
|
|
self.cloud_names = ['Lille1_1', 'Lille1_2', 'Lille2', 'Paris', 'ajaccio_2', 'ajaccio_57', 'dijon_9']
|
|
self.all_splits = [0, 1, 2, 3, 4, 5, 6]
|
|
self.validation_split = 1
|
|
# self.test_cloud_names = ['ajaccio_2', 'ajaccio_57', 'dijon_9']
|
|
self.test_splits = [4, 5, 6]
|
|
self.train_splits = [0, 2, 3]
|
|
|
|
# Number of models used per epoch
|
|
if self.set == 'training':
|
|
self.epoch_n = config.epoch_steps * config.batch_num
|
|
elif self.set in ['validation', 'test', 'ERF']:
|
|
self.epoch_n = config.validation_size * config.batch_num
|
|
else:
|
|
raise ValueError('Unknown set for NPM3D data: ', self.set)
|
|
|
|
# Stop data is not needed
|
|
if not load_data:
|
|
return
|
|
|
|
###################
|
|
# Prepare ply files
|
|
###################
|
|
|
|
self.prepare_NPM3D_ply()
|
|
|
|
################
|
|
# Load ply files
|
|
################
|
|
|
|
# List of training files
|
|
self.files = []
|
|
for i, f in enumerate(self.cloud_names):
|
|
if self.set == 'training':
|
|
if self.all_splits[i] in self.train_splits:
|
|
self.files += [join(ply_path, f + '.ply')]
|
|
elif self.set in ['validation', 'ERF']:
|
|
if self.all_splits[i] == self.validation_split:
|
|
self.files += [join(ply_path, f + '.ply')]
|
|
elif self.set == 'test':
|
|
if self.all_splits[i] in self.test_splits:
|
|
self.files += [join(ply_path, f + '.ply')]
|
|
else:
|
|
raise ValueError('Unknown set for NPM3D data: ', self.set)
|
|
print('The set is ' + str(self.set))
|
|
|
|
if self.set == 'training':
|
|
self.cloud_names = [f for i, f in enumerate(self.cloud_names)
|
|
if self.all_splits[i] in self.train_splits]
|
|
elif self.set in ['validation', 'ERF']:
|
|
self.cloud_names = [f for i, f in enumerate(self.cloud_names)
|
|
if self.all_splits[i] == self.validation_split]
|
|
elif self.set == 'test':
|
|
self.cloud_names = [f for i, f in enumerate(self.cloud_names)
|
|
if self.all_splits[i] in self.test_splits]
|
|
print('The files are ' + str(self.cloud_names))
|
|
|
|
if 0 < self.config.first_subsampling_dl <= 0.01:
|
|
raise ValueError('subsampling_parameter too low (should be over 1 cm')
|
|
|
|
# Initiate containers
|
|
self.input_trees = []
|
|
self.input_colors = []
|
|
self.input_labels = []
|
|
self.pot_trees = []
|
|
self.num_clouds = 0
|
|
self.test_proj = []
|
|
self.validation_labels = []
|
|
|
|
# Start loading
|
|
self.load_subsampled_clouds()
|
|
|
|
############################
|
|
# Batch selection parameters
|
|
############################
|
|
|
|
# Initialize value for batch limit (max number of points per batch).
|
|
self.batch_limit = torch.tensor([1], dtype=torch.float32)
|
|
self.batch_limit.share_memory_()
|
|
|
|
# Initialize potentials
|
|
if use_potentials:
|
|
self.potentials = []
|
|
self.min_potentials = []
|
|
self.argmin_potentials = []
|
|
for i, tree in enumerate(self.pot_trees):
|
|
self.potentials += [torch.from_numpy(np.random.rand(tree.data.shape[0]) * 1e-3)]
|
|
min_ind = int(torch.argmin(self.potentials[-1]))
|
|
self.argmin_potentials += [min_ind]
|
|
self.min_potentials += [float(self.potentials[-1][min_ind])]
|
|
|
|
# Share potential memory
|
|
self.argmin_potentials = torch.from_numpy(np.array(self.argmin_potentials, dtype=np.int64))
|
|
self.min_potentials = torch.from_numpy(np.array(self.min_potentials, dtype=np.float64))
|
|
self.argmin_potentials.share_memory_()
|
|
self.min_potentials.share_memory_()
|
|
for i, _ in enumerate(self.pot_trees):
|
|
self.potentials[i].share_memory_()
|
|
|
|
self.worker_waiting = torch.tensor([0 for _ in range(config.input_threads)], dtype=torch.int32)
|
|
self.worker_waiting.share_memory_()
|
|
self.epoch_inds = None
|
|
self.epoch_i = 0
|
|
|
|
else:
|
|
self.potentials = None
|
|
self.min_potentials = None
|
|
self.argmin_potentials = None
|
|
self.epoch_inds = torch.from_numpy(np.zeros((2, self.epoch_n), dtype=np.int64))
|
|
self.epoch_i = torch.from_numpy(np.zeros((1,), dtype=np.int64))
|
|
self.epoch_i.share_memory_()
|
|
self.epoch_inds.share_memory_()
|
|
|
|
self.worker_lock = Lock()
|
|
|
|
# For ERF visualization, we want only one cloud per batch and no randomness
|
|
if self.set == 'ERF':
|
|
self.batch_limit = torch.tensor([1], dtype=torch.float32)
|
|
self.batch_limit.share_memory_()
|
|
np.random.seed(42)
|
|
|
|
return
|
|
|
|
def __len__(self):
|
|
"""
|
|
Return the length of data here
|
|
"""
|
|
return len(self.cloud_names)
|
|
|
|
def __getitem__(self, batch_i):
|
|
"""
|
|
The main thread gives a list of indices to load a batch. Each worker is going to work in parallel to load a
|
|
different list of indices.
|
|
"""
|
|
|
|
if self.use_potentials:
|
|
return self.potential_item(batch_i)
|
|
else:
|
|
return self.random_item(batch_i)
|
|
|
|
def potential_item(self, batch_i, debug_workers=False):
|
|
|
|
t = [time.time()]
|
|
|
|
# Initiate concatanation lists
|
|
p_list = []
|
|
f_list = []
|
|
l_list = []
|
|
i_list = []
|
|
pi_list = []
|
|
ci_list = []
|
|
s_list = []
|
|
R_list = []
|
|
batch_n = 0
|
|
failed_attempts = 0
|
|
|
|
info = get_worker_info()
|
|
if info is not None:
|
|
wid = info.id
|
|
else:
|
|
wid = None
|
|
|
|
while True:
|
|
|
|
t += [time.time()]
|
|
|
|
if debug_workers:
|
|
message = ''
|
|
for wi in range(info.num_workers):
|
|
if wi == wid:
|
|
message += ' {:}X{:} '.format(bcolors.FAIL, bcolors.ENDC)
|
|
elif self.worker_waiting[wi] == 0:
|
|
message += ' '
|
|
elif self.worker_waiting[wi] == 1:
|
|
message += ' | '
|
|
elif self.worker_waiting[wi] == 2:
|
|
message += ' o '
|
|
print(message)
|
|
self.worker_waiting[wid] = 0
|
|
|
|
with self.worker_lock:
|
|
|
|
if debug_workers:
|
|
message = ''
|
|
for wi in range(info.num_workers):
|
|
if wi == wid:
|
|
message += ' {:}v{:} '.format(bcolors.OKGREEN, bcolors.ENDC)
|
|
elif self.worker_waiting[wi] == 0:
|
|
message += ' '
|
|
elif self.worker_waiting[wi] == 1:
|
|
message += ' | '
|
|
elif self.worker_waiting[wi] == 2:
|
|
message += ' o '
|
|
print(message)
|
|
self.worker_waiting[wid] = 1
|
|
|
|
# Get potential minimum
|
|
cloud_ind = int(torch.argmin(self.min_potentials))
|
|
point_ind = int(self.argmin_potentials[cloud_ind])
|
|
|
|
# Get potential points from tree structure
|
|
pot_points = np.array(self.pot_trees[cloud_ind].data, copy=False)
|
|
|
|
# Center point of input region
|
|
center_point = pot_points[point_ind, :].reshape(1, -1)
|
|
|
|
# Add a small noise to center point
|
|
if self.set != 'ERF':
|
|
center_point += np.random.normal(scale=self.config.in_radius / 10, size=center_point.shape)
|
|
|
|
# Indices of points in input region
|
|
pot_inds, dists = self.pot_trees[cloud_ind].query_radius(center_point,
|
|
r=self.config.in_radius,
|
|
return_distance=True)
|
|
|
|
d2s = np.square(dists[0])
|
|
pot_inds = pot_inds[0]
|
|
|
|
# Update potentials (Tukey weights)
|
|
if self.set != 'ERF':
|
|
tukeys = np.square(1 - d2s / np.square(self.config.in_radius))
|
|
tukeys[d2s > np.square(self.config.in_radius)] = 0
|
|
self.potentials[cloud_ind][pot_inds] += tukeys
|
|
min_ind = torch.argmin(self.potentials[cloud_ind])
|
|
self.min_potentials[[cloud_ind]] = self.potentials[cloud_ind][min_ind]
|
|
self.argmin_potentials[[cloud_ind]] = min_ind
|
|
|
|
t += [time.time()]
|
|
|
|
# Get points from tree structure
|
|
points = np.array(self.input_trees[cloud_ind].data, copy=False)
|
|
|
|
# Indices of points in input region
|
|
input_inds = self.input_trees[cloud_ind].query_radius(center_point,
|
|
r=self.config.in_radius)[0]
|
|
|
|
t += [time.time()]
|
|
|
|
# Number collected
|
|
n = input_inds.shape[0]
|
|
|
|
# Safe check for empty spheres
|
|
if n < 2:
|
|
failed_attempts += 1
|
|
if failed_attempts > 100 * self.config.batch_num:
|
|
raise ValueError('It seems this dataset only containes empty input spheres')
|
|
t += [time.time()]
|
|
t += [time.time()]
|
|
continue
|
|
|
|
# Collect labels and colors
|
|
input_points = (points[input_inds] - center_point).astype(np.float32)
|
|
# input_colors = self.input_colors[cloud_ind][input_inds]
|
|
if self.set in ['test', 'ERF']:
|
|
input_labels = np.zeros(input_points.shape[0])
|
|
else:
|
|
input_labels = self.input_labels[cloud_ind][input_inds]
|
|
input_labels = np.array([self.label_to_idx[l] for l in input_labels])
|
|
|
|
t += [time.time()]
|
|
|
|
# Data augmentation
|
|
input_points, scale, R = self.augmentation_transform(input_points)
|
|
|
|
# Color augmentation
|
|
# if np.random.rand() > self.config.augment_color:
|
|
# input_colors *= 0
|
|
|
|
# Get original height as additional feature
|
|
# input_features = np.hstack((input_colors, input_points[:, 2:] + center_point[:, 2:])).astype(np.float32)
|
|
input_features = np.hstack((input_points[:, 2:] + center_point[:, 2:])).astype(np.float32)
|
|
|
|
t += [time.time()]
|
|
|
|
# Stack batch
|
|
p_list += [input_points]
|
|
f_list += [input_features]
|
|
l_list += [input_labels]
|
|
pi_list += [input_inds]
|
|
i_list += [point_ind]
|
|
ci_list += [cloud_ind]
|
|
s_list += [scale]
|
|
R_list += [R]
|
|
|
|
# Update batch size
|
|
batch_n += n
|
|
|
|
# In case batch is full, stop
|
|
if batch_n > int(self.batch_limit):
|
|
break
|
|
|
|
# Randomly drop some points (act as an augmentation process and a safety for GPU memory consumption)
|
|
# if n > int(self.batch_limit):
|
|
# input_inds = np.random.choice(input_inds, size=int(self.batch_limit) - 1, replace=False)
|
|
# n = input_inds.shape[0]
|
|
|
|
###################
|
|
# Concatenate batch
|
|
###################
|
|
|
|
stacked_points = np.concatenate(p_list, axis=0)
|
|
features = np.concatenate(f_list, axis=0)
|
|
labels = np.concatenate(l_list, axis=0)
|
|
point_inds = np.array(i_list, dtype=np.int32)
|
|
cloud_inds = np.array(ci_list, dtype=np.int32)
|
|
input_inds = np.concatenate(pi_list, axis=0)
|
|
stack_lengths = np.array([pp.shape[0] for pp in p_list], dtype=np.int32)
|
|
scales = np.array(s_list, dtype=np.float32)
|
|
rots = np.stack(R_list, axis=0)
|
|
|
|
# Input features
|
|
stacked_features = np.ones_like(stacked_points[:, :1], dtype=np.float32)
|
|
if self.config.in_features_dim == 1:
|
|
pass
|
|
elif self.config.in_features_dim == 4:
|
|
stacked_features = np.hstack((stacked_features, features[:, :3]))
|
|
elif self.config.in_features_dim == 5:
|
|
stacked_features = np.hstack((stacked_features, features))
|
|
else:
|
|
raise ValueError('Only accepted input dimensions are 1, 4 and 7 (without and with XYZ)')
|
|
|
|
#######################
|
|
# Create network inputs
|
|
#######################
|
|
#
|
|
# Points, neighbors, pooling indices for each layers
|
|
#
|
|
|
|
t += [time.time()]
|
|
|
|
# Get the whole input list
|
|
input_list = self.segmentation_inputs(stacked_points,
|
|
stacked_features,
|
|
labels,
|
|
stack_lengths)
|
|
|
|
t += [time.time()]
|
|
|
|
# Add scale and rotation for testing
|
|
input_list += [scales, rots, cloud_inds, point_inds, input_inds]
|
|
|
|
if debug_workers:
|
|
message = ''
|
|
for wi in range(info.num_workers):
|
|
if wi == wid:
|
|
message += ' {:}0{:} '.format(bcolors.OKBLUE, bcolors.ENDC)
|
|
elif self.worker_waiting[wi] == 0:
|
|
message += ' '
|
|
elif self.worker_waiting[wi] == 1:
|
|
message += ' | '
|
|
elif self.worker_waiting[wi] == 2:
|
|
message += ' o '
|
|
print(message)
|
|
self.worker_waiting[wid] = 2
|
|
|
|
t += [time.time()]
|
|
|
|
# Display timings
|
|
debugT = False
|
|
if debugT:
|
|
print('\n************************\n')
|
|
print('Timings:')
|
|
ti = 0
|
|
N = 5
|
|
mess = 'Init ...... {:5.1f}ms /'
|
|
loop_times = [1000 * (t[ti + N * i + 1] - t[ti + N * i]) for i in range(len(stack_lengths))]
|
|
for dt in loop_times:
|
|
mess += ' {:5.1f}'.format(dt)
|
|
print(mess.format(np.sum(loop_times)))
|
|
ti += 1
|
|
mess = 'Pots ...... {:5.1f}ms /'
|
|
loop_times = [1000 * (t[ti + N * i + 1] - t[ti + N * i]) for i in range(len(stack_lengths))]
|
|
for dt in loop_times:
|
|
mess += ' {:5.1f}'.format(dt)
|
|
print(mess.format(np.sum(loop_times)))
|
|
ti += 1
|
|
mess = 'Sphere .... {:5.1f}ms /'
|
|
loop_times = [1000 * (t[ti + N * i + 1] - t[ti + N * i]) for i in range(len(stack_lengths))]
|
|
for dt in loop_times:
|
|
mess += ' {:5.1f}'.format(dt)
|
|
print(mess.format(np.sum(loop_times)))
|
|
ti += 1
|
|
mess = 'Collect ... {:5.1f}ms /'
|
|
loop_times = [1000 * (t[ti + N * i + 1] - t[ti + N * i]) for i in range(len(stack_lengths))]
|
|
for dt in loop_times:
|
|
mess += ' {:5.1f}'.format(dt)
|
|
print(mess.format(np.sum(loop_times)))
|
|
ti += 1
|
|
mess = 'Augment ... {:5.1f}ms /'
|
|
loop_times = [1000 * (t[ti + N * i + 1] - t[ti + N * i]) for i in range(len(stack_lengths))]
|
|
for dt in loop_times:
|
|
mess += ' {:5.1f}'.format(dt)
|
|
print(mess.format(np.sum(loop_times)))
|
|
ti += N * (len(stack_lengths) - 1) + 1
|
|
print('concat .... {:5.1f}ms'.format(1000 * (t[ti + 1] - t[ti])))
|
|
ti += 1
|
|
print('input ..... {:5.1f}ms'.format(1000 * (t[ti + 1] - t[ti])))
|
|
ti += 1
|
|
print('stack ..... {:5.1f}ms'.format(1000 * (t[ti + 1] - t[ti])))
|
|
ti += 1
|
|
print('\n************************\n')
|
|
return input_list
|
|
|
|
def random_item(self, batch_i):
|
|
|
|
# Initiate concatanation lists
|
|
p_list = []
|
|
f_list = []
|
|
l_list = []
|
|
i_list = []
|
|
pi_list = []
|
|
ci_list = []
|
|
s_list = []
|
|
R_list = []
|
|
batch_n = 0
|
|
failed_attempts = 0
|
|
|
|
while True:
|
|
|
|
with self.worker_lock:
|
|
|
|
# Get potential minimum
|
|
cloud_ind = int(self.epoch_inds[0, self.epoch_i])
|
|
point_ind = int(self.epoch_inds[1, self.epoch_i])
|
|
|
|
# Update epoch indice
|
|
self.epoch_i += 1
|
|
if self.epoch_i >= int(self.epoch_inds.shape[1]):
|
|
self.epoch_i -= int(self.epoch_inds.shape[1])
|
|
|
|
# Get points from tree structure
|
|
points = np.array(self.input_trees[cloud_ind].data, copy=False)
|
|
|
|
# Center point of input region
|
|
center_point = points[point_ind, :].reshape(1, -1)
|
|
|
|
# Add a small noise to center point
|
|
if self.set != 'ERF':
|
|
center_point += np.random.normal(scale=self.config.in_radius / 10, size=center_point.shape)
|
|
|
|
# Indices of points in input region
|
|
input_inds = self.input_trees[cloud_ind].query_radius(center_point,
|
|
r=self.config.in_radius)[0]
|
|
|
|
# Number collected
|
|
n = input_inds.shape[0]
|
|
|
|
# Safe check for empty spheres
|
|
if n < 2:
|
|
failed_attempts += 1
|
|
if failed_attempts > 100 * self.config.batch_num:
|
|
raise ValueError('It seems this dataset only containes empty input spheres')
|
|
continue
|
|
|
|
# Collect labels and colors
|
|
input_points = (points[input_inds] - center_point).astype(np.float32)
|
|
# input_colors = self.input_colors[cloud_ind][input_inds]
|
|
if self.set in ['test', 'ERF']:
|
|
input_labels = np.zeros(input_points.shape[0])
|
|
else:
|
|
input_labels = self.input_labels[cloud_ind][input_inds]
|
|
input_labels = np.array([self.label_to_idx[l] for l in input_labels])
|
|
|
|
# Data augmentation
|
|
input_points, scale, R = self.augmentation_transform(input_points)
|
|
|
|
# Color augmentation
|
|
# if np.random.rand() > self.config.augment_color:
|
|
# input_colors *= 0
|
|
|
|
# Get original height as additional feature
|
|
input_features = np.hstack((input_colors, input_points[:, 2:] + center_point[:, 2:])).astype(np.float32)
|
|
|
|
# Stack batch
|
|
p_list += [input_points]
|
|
f_list += [input_features]
|
|
l_list += [input_labels]
|
|
pi_list += [input_inds]
|
|
i_list += [point_ind]
|
|
ci_list += [cloud_ind]
|
|
s_list += [scale]
|
|
R_list += [R]
|
|
|
|
# Update batch size
|
|
batch_n += n
|
|
|
|
# In case batch is full, stop
|
|
if batch_n > int(self.batch_limit):
|
|
break
|
|
|
|
# Randomly drop some points (act as an augmentation process and a safety for GPU memory consumption)
|
|
# if n > int(self.batch_limit):
|
|
# input_inds = np.random.choice(input_inds, size=int(self.batch_limit) - 1, replace=False)
|
|
# n = input_inds.shape[0]
|
|
|
|
###################
|
|
# Concatenate batch
|
|
###################
|
|
|
|
stacked_points = np.concatenate(p_list, axis=0)
|
|
features = np.concatenate(f_list, axis=0)
|
|
labels = np.concatenate(l_list, axis=0)
|
|
point_inds = np.array(i_list, dtype=np.int32)
|
|
cloud_inds = np.array(ci_list, dtype=np.int32)
|
|
input_inds = np.concatenate(pi_list, axis=0)
|
|
stack_lengths = np.array([pp.shape[0] for pp in p_list], dtype=np.int32)
|
|
scales = np.array(s_list, dtype=np.float32)
|
|
rots = np.stack(R_list, axis=0)
|
|
|
|
# Input features
|
|
stacked_features = np.ones_like(stacked_points[:, :1], dtype=np.float32)
|
|
if self.config.in_features_dim == 1:
|
|
pass
|
|
elif self.config.in_features_dim == 4:
|
|
stacked_features = np.hstack((stacked_features, features[:, :3]))
|
|
elif self.config.in_features_dim == 5:
|
|
stacked_features = np.hstack((stacked_features, features))
|
|
else:
|
|
raise ValueError('Only accepted input dimensions are 1, 4 and 7 (without and with XYZ)')
|
|
|
|
#######################
|
|
# Create network inputs
|
|
#######################
|
|
#
|
|
# Points, neighbors, pooling indices for each layers
|
|
#
|
|
|
|
# Get the whole input list
|
|
input_list = self.segmentation_inputs(stacked_points,
|
|
stacked_features,
|
|
labels,
|
|
stack_lengths)
|
|
|
|
# Add scale and rotation for testing
|
|
input_list += [scales, rots, cloud_inds, point_inds, input_inds]
|
|
|
|
return input_list
|
|
|
|
def prepare_NPM3D_ply(self):
|
|
|
|
print('\nPreparing ply files')
|
|
t0 = time.time()
|
|
|
|
# Folder for the ply files
|
|
ply_path = join(self.path, self.train_path)
|
|
if not exists(ply_path):
|
|
makedirs(ply_path)
|
|
|
|
for cloud_name in self.cloud_names:
|
|
|
|
# Pass if the cloud has already been computed
|
|
cloud_file = join(ply_path, cloud_name + '.ply')
|
|
if exists(cloud_file):
|
|
continue
|
|
|
|
original_ply = read_ply(join(self.path, self.original_ply_path, cloud_name + '.ply'))
|
|
|
|
# Initiate containers
|
|
cloud_x = original_ply['x']
|
|
cloud_y = original_ply['y']
|
|
cloud_z = original_ply['z']
|
|
cloud_x = cloud_x - (cloud_x.min())
|
|
cloud_y = cloud_y - (cloud_y.min())
|
|
cloud_z = cloud_z - (cloud_z.min())
|
|
|
|
# Reshape
|
|
cloud_x = cloud_x.reshape(len(cloud_x), 1)
|
|
cloud_y = cloud_y.reshape(len(cloud_y), 1)
|
|
cloud_z = cloud_z.reshape(len(cloud_z), 1)
|
|
|
|
# Astype
|
|
cloud_x = cloud_x.astype(np.float32)
|
|
cloud_y = cloud_y.astype(np.float32)
|
|
cloud_z = cloud_z.astype(np.float32)
|
|
|
|
# Stack
|
|
cloud_points = np.hstack((cloud_x, cloud_y, cloud_z))
|
|
|
|
# Labels
|
|
if cloud_name in ['ajaccio_2', 'ajaccio_57', 'dijon_9']:
|
|
|
|
field_names = ['x', 'y', 'z']
|
|
write_ply(join(ply_path, cloud_name + '.ply'), cloud_points, field_names)
|
|
|
|
else:
|
|
labels = original_ply['class']
|
|
labels = labels.astype(np.int32)
|
|
labels = labels.reshape(len(labels), 1)
|
|
|
|
# Save as ply
|
|
field_names = ['x', 'y', 'z', 'class']
|
|
write_ply(join(ply_path, cloud_name + '.ply'), [cloud_points, labels], field_names)
|
|
|
|
print('Done in {:.1f}s'.format(time.time() - t0))
|
|
return
|
|
|
|
def load_subsampled_clouds(self):
|
|
|
|
# Parameter
|
|
dl = self.config.first_subsampling_dl
|
|
|
|
# Create path for files
|
|
tree_path = join(self.path, 'input_{:.3f}'.format(dl))
|
|
if not exists(tree_path):
|
|
makedirs(tree_path)
|
|
|
|
##############
|
|
# Load KDTrees
|
|
##############
|
|
|
|
for i, file_path in enumerate(self.files):
|
|
|
|
# Restart timer
|
|
t0 = time.time()
|
|
|
|
# Get cloud name
|
|
cloud_name = self.cloud_names[i]
|
|
|
|
# Name of the input files
|
|
KDTree_file = join(tree_path, '{:s}_KDTree.pkl'.format(cloud_name))
|
|
sub_ply_file = join(tree_path, '{:s}.ply'.format(cloud_name))
|
|
|
|
# Check if inputs have already been computed
|
|
if exists(KDTree_file):
|
|
print('\nFound KDTree for cloud {:s}, subsampled at {:.3f}'.format(cloud_name, dl))
|
|
|
|
# read ply with data
|
|
data = read_ply(sub_ply_file)
|
|
# sub_colors = np.vstack((data['red'], data['green'], data['blue'])).T
|
|
sub_labels = data['class']
|
|
|
|
# Read pkl with search tree
|
|
with open(KDTree_file, 'rb') as f:
|
|
search_tree = pickle.load(f)
|
|
|
|
else:
|
|
print('\nPreparing KDTree for cloud {:s}, subsampled at {:.3f}'.format(cloud_name, dl))
|
|
|
|
# Read ply file
|
|
data = read_ply(file_path)
|
|
points = np.vstack((data['x'], data['y'], data['z'])).T
|
|
# colors = np.vstack((data['red'], data['green'], data['blue'])).T
|
|
|
|
# Fake labels for test data
|
|
if self.set == 'test':
|
|
labels = np.zeros((data.shape[0],), dtype=np.int32)
|
|
else:
|
|
labels = data['class']
|
|
|
|
# Subsample cloud
|
|
sub_points, sub_labels = grid_subsampling(points,
|
|
labels=labels,
|
|
sampleDl=dl)
|
|
|
|
# Rescale float color and squeeze label
|
|
# sub_colors = sub_colors / 255
|
|
sub_labels = np.squeeze(sub_labels)
|
|
|
|
# Get chosen neighborhoods
|
|
search_tree = KDTree(sub_points, leaf_size=10)
|
|
# search_tree = nnfln.KDTree(n_neighbors=1, metric='L2', leaf_size=10)
|
|
# search_tree.fit(sub_points)
|
|
|
|
# Save KDTree
|
|
with open(KDTree_file, 'wb') as f:
|
|
pickle.dump(search_tree, f)
|
|
|
|
# Save ply
|
|
write_ply(sub_ply_file,
|
|
[sub_points, sub_labels],
|
|
['x', 'y', 'z', 'class'])
|
|
|
|
# Fill data containers
|
|
self.input_trees += [search_tree]
|
|
# self.input_colors += [sub_colors]
|
|
self.input_labels += [sub_labels]
|
|
|
|
size = sub_labels.shape[0] * 4 * 7
|
|
print('{:.1f} MB loaded in {:.1f}s'.format(size * 1e-6, time.time() - t0))
|
|
|
|
############################
|
|
# Coarse potential locations
|
|
############################
|
|
|
|
# Only necessary for validation and test sets
|
|
if self.use_potentials:
|
|
print('\nPreparing potentials')
|
|
|
|
# Restart timer
|
|
t0 = time.time()
|
|
|
|
pot_dl = self.config.in_radius / 10
|
|
cloud_ind = 0
|
|
|
|
for i, file_path in enumerate(self.files):
|
|
|
|
# Get cloud name
|
|
cloud_name = self.cloud_names[i]
|
|
|
|
# Name of the input files
|
|
coarse_KDTree_file = join(tree_path, '{:s}_coarse_KDTree.pkl'.format(cloud_name))
|
|
|
|
# Check if inputs have already been computed
|
|
if exists(coarse_KDTree_file):
|
|
# Read pkl with search tree
|
|
with open(coarse_KDTree_file, 'rb') as f:
|
|
search_tree = pickle.load(f)
|
|
|
|
else:
|
|
# Subsample cloud
|
|
sub_points = np.array(self.input_trees[cloud_ind].data, copy=False)
|
|
coarse_points = grid_subsampling(sub_points.astype(np.float32), sampleDl=pot_dl)
|
|
|
|
# Get chosen neighborhoods
|
|
search_tree = KDTree(coarse_points, leaf_size=10)
|
|
|
|
# Save KDTree
|
|
with open(coarse_KDTree_file, 'wb') as f:
|
|
pickle.dump(search_tree, f)
|
|
|
|
# Fill data containers
|
|
self.pot_trees += [search_tree]
|
|
cloud_ind += 1
|
|
|
|
print('Done in {:.1f}s'.format(time.time() - t0))
|
|
|
|
######################
|
|
# Reprojection indices
|
|
######################
|
|
|
|
# Get number of clouds
|
|
self.num_clouds = len(self.input_trees)
|
|
|
|
# Only necessary for validation and test sets
|
|
if self.set in ['validation', 'test']:
|
|
|
|
print('\nPreparing reprojection indices for testing')
|
|
|
|
# Get validation/test reprojection indices
|
|
for i, file_path in enumerate(self.files):
|
|
|
|
# Restart timer
|
|
t0 = time.time()
|
|
|
|
# Get info on this cloud
|
|
cloud_name = self.cloud_names[i]
|
|
|
|
# File name for saving
|
|
proj_file = join(tree_path, '{:s}_proj.pkl'.format(cloud_name))
|
|
|
|
# Try to load previous indices
|
|
if exists(proj_file):
|
|
with open(proj_file, 'rb') as f:
|
|
proj_inds, labels = pickle.load(f)
|
|
else:
|
|
data = read_ply(file_path)
|
|
points = np.vstack((data['x'], data['y'], data['z'])).T
|
|
|
|
# Fake labels
|
|
if self.set == 'test':
|
|
labels = np.zeros((data.shape[0],), dtype=np.int32)
|
|
else:
|
|
labels = data['class']
|
|
|
|
# Compute projection inds
|
|
idxs = self.input_trees[i].query(points, return_distance=False)
|
|
# dists, idxs = self.input_trees[i_cloud].kneighbors(points)
|
|
proj_inds = np.squeeze(idxs).astype(np.int32)
|
|
|
|
# Save
|
|
with open(proj_file, 'wb') as f:
|
|
pickle.dump([proj_inds, labels], f)
|
|
|
|
self.test_proj += [proj_inds]
|
|
self.validation_labels += [labels]
|
|
print('{:s} done in {:.1f}s'.format(cloud_name, time.time() - t0))
|
|
|
|
print()
|
|
return
|
|
|
|
def load_evaluation_points(self, file_path):
|
|
"""
|
|
Load points (from test or validation split) on which the metrics should be evaluated
|
|
"""
|
|
|
|
# Get original points
|
|
data = read_ply(file_path)
|
|
return np.vstack((data['x'], data['y'], data['z'])).T
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Utility classes definition
|
|
# \********************************/
|
|
|
|
|
|
class NPM3DSampler(Sampler):
|
|
"""Sampler for NPM3D"""
|
|
|
|
def __init__(self, dataset: NPM3DDataset):
|
|
Sampler.__init__(self, dataset)
|
|
|
|
# Dataset used by the sampler (no copy is made in memory)
|
|
self.dataset = dataset
|
|
|
|
# Number of step per epoch
|
|
if dataset.set == 'training':
|
|
self.N = dataset.config.epoch_steps
|
|
else:
|
|
self.N = dataset.config.validation_size
|
|
|
|
return
|
|
|
|
def __iter__(self):
|
|
"""
|
|
Yield next batch indices here. In this dataset, this is a dummy sampler that yield the index of batch element
|
|
(input sphere) in epoch instead of the list of point indices
|
|
"""
|
|
|
|
if not self.dataset.use_potentials:
|
|
|
|
# Initiate current epoch ind
|
|
self.dataset.epoch_i *= 0
|
|
self.dataset.epoch_inds *= 0
|
|
|
|
# Initiate container for indices
|
|
all_epoch_inds = np.zeros((2, 0), dtype=np.int64)
|
|
|
|
# Number of sphere centers taken per class in each cloud
|
|
num_centers = self.N * self.dataset.config.batch_num
|
|
random_pick_n = int(np.ceil(num_centers / self.dataset.config.num_classes))
|
|
|
|
# Choose random points of each class for each cloud
|
|
epoch_indices = np.zeros((2, 0), dtype=np.int64)
|
|
for label_ind, label in enumerate(self.dataset.label_values):
|
|
if label not in self.dataset.ignored_labels:
|
|
|
|
# Gather indices of the points with this label in all the input clouds
|
|
all_label_indices = []
|
|
for cloud_ind, cloud_labels in enumerate(self.dataset.input_labels):
|
|
label_indices = np.where(np.equal(cloud_labels, label))[0]
|
|
all_label_indices.append(
|
|
np.vstack((np.full(label_indices.shape, cloud_ind, dtype=np.int64), label_indices)))
|
|
|
|
# Stack them: [2, N1+N2+...]
|
|
all_label_indices = np.hstack(all_label_indices)
|
|
|
|
# Select a a random number amongst them
|
|
N_inds = all_label_indices.shape[1]
|
|
if N_inds < random_pick_n:
|
|
chosen_label_inds = np.zeros((2, 0), dtype=np.int64)
|
|
while chosen_label_inds.shape[1] < random_pick_n:
|
|
chosen_label_inds = np.hstack(
|
|
(chosen_label_inds, all_label_indices[:, np.random.permutation(N_inds)]))
|
|
warnings.warn('When choosing random epoch indices (use_potentials=False), \
|
|
class {:d}: {:s} only had {:d} available points, while we \
|
|
needed {:d}. Repeating indices in the same epoch'.format(label,
|
|
self.dataset.label_names[
|
|
label_ind],
|
|
N_inds,
|
|
random_pick_n))
|
|
|
|
elif N_inds < 50 * random_pick_n:
|
|
rand_inds = np.random.choice(N_inds, size=random_pick_n, replace=False)
|
|
chosen_label_inds = all_label_indices[:, rand_inds]
|
|
|
|
else:
|
|
chosen_label_inds = np.zeros((2, 0), dtype=np.int64)
|
|
while chosen_label_inds.shape[1] < random_pick_n:
|
|
rand_inds = np.unique(np.random.choice(N_inds, size=2 * random_pick_n, replace=True))
|
|
chosen_label_inds = np.hstack((chosen_label_inds, all_label_indices[:, rand_inds]))
|
|
chosen_label_inds = chosen_label_inds[:, :random_pick_n]
|
|
|
|
# Stack for each label
|
|
all_epoch_inds = np.hstack((all_epoch_inds, chosen_label_inds))
|
|
|
|
# Random permutation of the indices
|
|
random_order = np.random.permutation(all_epoch_inds.shape[1])[:num_centers]
|
|
all_epoch_inds = all_epoch_inds[:, random_order].astype(np.int64)
|
|
|
|
# Update epoch inds
|
|
self.dataset.epoch_inds += torch.from_numpy(all_epoch_inds)
|
|
|
|
# Generator loop
|
|
for i in range(self.N):
|
|
yield i
|
|
|
|
def __len__(self):
|
|
"""
|
|
The number of yielded samples is variable
|
|
"""
|
|
return self.N
|
|
|
|
def fast_calib(self):
|
|
"""
|
|
This method calibrates the batch sizes while ensuring the potentials are well initialized. Indeed on a dataset
|
|
like Semantic3D, before potential have been updated over the dataset, there are cahnces that all the dense area
|
|
are picked in the begining and in the end, we will have very large batch of small point clouds
|
|
:return:
|
|
"""
|
|
|
|
# Estimated average batch size and target value
|
|
estim_b = 0
|
|
target_b = self.dataset.config.batch_num
|
|
|
|
# Calibration parameters
|
|
low_pass_T = 10
|
|
Kp = 100.0
|
|
finer = False
|
|
breaking = False
|
|
|
|
# Convergence parameters
|
|
smooth_errors = []
|
|
converge_threshold = 0.1
|
|
|
|
t = [time.time()]
|
|
last_display = time.time()
|
|
mean_dt = np.zeros(2)
|
|
|
|
for epoch in range(10):
|
|
for i, test in enumerate(self):
|
|
|
|
# New time
|
|
t = t[-1:]
|
|
t += [time.time()]
|
|
|
|
# batch length
|
|
b = len(test)
|
|
|
|
# Update estim_b (low pass filter)
|
|
estim_b += (b - estim_b) / low_pass_T
|
|
|
|
# Estimate error (noisy)
|
|
error = target_b - b
|
|
|
|
# Save smooth errors for convergene check
|
|
smooth_errors.append(target_b - estim_b)
|
|
if len(smooth_errors) > 10:
|
|
smooth_errors = smooth_errors[1:]
|
|
|
|
# Update batch limit with P controller
|
|
self.dataset.batch_limit += Kp * error
|
|
|
|
# finer low pass filter when closing in
|
|
if not finer and np.abs(estim_b - target_b) < 1:
|
|
low_pass_T = 100
|
|
finer = True
|
|
|
|
# Convergence
|
|
if finer and np.max(np.abs(smooth_errors)) < converge_threshold:
|
|
breaking = True
|
|
break
|
|
|
|
# Average timing
|
|
t += [time.time()]
|
|
mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1]))
|
|
|
|
# Console display (only one per second)
|
|
if (t[-1] - last_display) > 1.0:
|
|
last_display = t[-1]
|
|
message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}, // {:.1f}ms {:.1f}ms'
|
|
print(message.format(i,
|
|
estim_b,
|
|
int(self.dataset.batch_limit),
|
|
1000 * mean_dt[0],
|
|
1000 * mean_dt[1]))
|
|
|
|
if breaking:
|
|
break
|
|
|
|
def calibration(self, dataloader, untouched_ratio=0.9, verbose=False, force_redo=False):
|
|
"""
|
|
Method performing batch and neighbors calibration.
|
|
Batch calibration: Set "batch_limit" (the maximum number of points allowed in every batch) so that the
|
|
average batch size (number of stacked pointclouds) is the one asked.
|
|
Neighbors calibration: Set the "neighborhood_limits" (the maximum number of neighbors allowed in convolutions)
|
|
so that 90% of the neighborhoods remain untouched. There is a limit for each layer.
|
|
"""
|
|
|
|
##############################
|
|
# Previously saved calibration
|
|
##############################
|
|
|
|
print('\nStarting Calibration (use verbose=True for more details)')
|
|
t0 = time.time()
|
|
|
|
redo = force_redo
|
|
|
|
# Batch limit
|
|
# ***********
|
|
|
|
# Load batch_limit dictionary
|
|
batch_lim_file = join(self.dataset.path, 'batch_limits.pkl')
|
|
if exists(batch_lim_file):
|
|
with open(batch_lim_file, 'rb') as file:
|
|
batch_lim_dict = pickle.load(file)
|
|
else:
|
|
batch_lim_dict = {}
|
|
|
|
# Check if the batch limit associated with current parameters exists
|
|
if self.dataset.use_potentials:
|
|
sampler_method = 'potentials'
|
|
else:
|
|
sampler_method = 'random'
|
|
key = '{:s}_{:.3f}_{:.3f}_{:d}'.format(sampler_method,
|
|
self.dataset.config.in_radius,
|
|
self.dataset.config.first_subsampling_dl,
|
|
self.dataset.config.batch_num)
|
|
if not redo and key in batch_lim_dict:
|
|
self.dataset.batch_limit[0] = batch_lim_dict[key]
|
|
else:
|
|
redo = True
|
|
|
|
if verbose:
|
|
print('\nPrevious calibration found:')
|
|
print('Check batch limit dictionary')
|
|
if key in batch_lim_dict:
|
|
color = bcolors.OKGREEN
|
|
v = str(int(batch_lim_dict[key]))
|
|
else:
|
|
color = bcolors.FAIL
|
|
v = '?'
|
|
print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC))
|
|
|
|
# Neighbors limit
|
|
# ***************
|
|
|
|
# Load neighb_limits dictionary
|
|
neighb_lim_file = join(self.dataset.path, 'neighbors_limits.pkl')
|
|
if exists(neighb_lim_file):
|
|
with open(neighb_lim_file, 'rb') as file:
|
|
neighb_lim_dict = pickle.load(file)
|
|
else:
|
|
neighb_lim_dict = {}
|
|
|
|
# Check if the limit associated with current parameters exists (for each layer)
|
|
neighb_limits = []
|
|
for layer_ind in range(self.dataset.config.num_layers):
|
|
|
|
dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind)
|
|
if self.dataset.config.deform_layers[layer_ind]:
|
|
r = dl * self.dataset.config.deform_radius
|
|
else:
|
|
r = dl * self.dataset.config.conv_radius
|
|
|
|
key = '{:.3f}_{:.3f}'.format(dl, r)
|
|
if key in neighb_lim_dict:
|
|
neighb_limits += [neighb_lim_dict[key]]
|
|
|
|
if not redo and len(neighb_limits) == self.dataset.config.num_layers:
|
|
self.dataset.neighborhood_limits = neighb_limits
|
|
else:
|
|
redo = True
|
|
|
|
if verbose:
|
|
print('Check neighbors limit dictionary')
|
|
for layer_ind in range(self.dataset.config.num_layers):
|
|
dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind)
|
|
if self.dataset.config.deform_layers[layer_ind]:
|
|
r = dl * self.dataset.config.deform_radius
|
|
else:
|
|
r = dl * self.dataset.config.conv_radius
|
|
key = '{:.3f}_{:.3f}'.format(dl, r)
|
|
|
|
if key in neighb_lim_dict:
|
|
color = bcolors.OKGREEN
|
|
v = str(neighb_lim_dict[key])
|
|
else:
|
|
color = bcolors.FAIL
|
|
v = '?'
|
|
print('{:}\"{:s}\": {:s}{:}'.format(color, key, v, bcolors.ENDC))
|
|
|
|
if redo:
|
|
|
|
############################
|
|
# Neighbors calib parameters
|
|
############################
|
|
|
|
# From config parameter, compute higher bound of neighbors number in a neighborhood
|
|
hist_n = int(np.ceil(4 / 3 * np.pi * (self.dataset.config.deform_radius + 1) ** 3))
|
|
|
|
# Histogram of neighborhood sizes
|
|
neighb_hists = np.zeros((self.dataset.config.num_layers, hist_n), dtype=np.int32)
|
|
|
|
########################
|
|
# Batch calib parameters
|
|
########################
|
|
|
|
# Estimated average batch size and target value
|
|
estim_b = 0
|
|
target_b = self.dataset.config.batch_num
|
|
|
|
# Expected batch size order of magnitude
|
|
expected_N = 100000
|
|
|
|
# Calibration parameters. Higher means faster but can also become unstable
|
|
# Reduce Kp and Kd if your GP Uis small as the total number of points per batch will be smaller
|
|
low_pass_T = 100
|
|
Kp = expected_N / 200
|
|
Ki = 0.001 * Kp
|
|
Kd = 5 * Kp
|
|
finer = False
|
|
stabilized = False
|
|
|
|
# Convergence parameters
|
|
smooth_errors = []
|
|
converge_threshold = 0.1
|
|
|
|
# Loop parameters
|
|
last_display = time.time()
|
|
i = 0
|
|
breaking = False
|
|
error_I = 0
|
|
error_D = 0
|
|
last_error = 0
|
|
|
|
debug_in = []
|
|
debug_out = []
|
|
debug_b = []
|
|
debug_estim_b = []
|
|
|
|
#####################
|
|
# Perform calibration
|
|
#####################
|
|
|
|
# number of batch per epoch
|
|
sample_batches = 999
|
|
for epoch in range((sample_batches // self.N) + 1):
|
|
for batch_i, batch in enumerate(dataloader):
|
|
|
|
# Update neighborhood histogram
|
|
counts = [np.sum(neighb_mat.numpy() < neighb_mat.shape[0], axis=1) for neighb_mat in
|
|
batch.neighbors]
|
|
hists = [np.bincount(c, minlength=hist_n)[:hist_n] for c in counts]
|
|
neighb_hists += np.vstack(hists)
|
|
|
|
# batch length
|
|
b = len(batch.cloud_inds)
|
|
|
|
# Update estim_b (low pass filter)
|
|
estim_b += (b - estim_b) / low_pass_T
|
|
|
|
# Estimate error (noisy)
|
|
error = target_b - b
|
|
error_I += error
|
|
error_D = error - last_error
|
|
last_error = error
|
|
|
|
# Save smooth errors for convergene check
|
|
smooth_errors.append(target_b - estim_b)
|
|
if len(smooth_errors) > 30:
|
|
smooth_errors = smooth_errors[1:]
|
|
|
|
# Update batch limit with P controller
|
|
self.dataset.batch_limit += Kp * error + Ki * error_I + Kd * error_D
|
|
|
|
# Unstability detection
|
|
if not stabilized and self.dataset.batch_limit < 0:
|
|
Kp *= 0.1
|
|
Ki *= 0.1
|
|
Kd *= 0.1
|
|
stabilized = True
|
|
|
|
# finer low pass filter when closing in
|
|
if not finer and np.abs(estim_b - target_b) < 1:
|
|
low_pass_T = 100
|
|
finer = True
|
|
|
|
# Convergence
|
|
if finer and np.max(np.abs(smooth_errors)) < converge_threshold:
|
|
breaking = True
|
|
break
|
|
|
|
i += 1
|
|
t = time.time()
|
|
|
|
# Console display (only one per second)
|
|
if verbose and (t - last_display) > 1.0:
|
|
last_display = t
|
|
message = 'Step {:5d} estim_b ={:5.2f} batch_limit ={:7d}'
|
|
print(message.format(i,
|
|
estim_b,
|
|
int(self.dataset.batch_limit)))
|
|
|
|
# Debug plots
|
|
debug_in.append(int(batch.points[0].shape[0]))
|
|
debug_out.append(int(self.dataset.batch_limit))
|
|
debug_b.append(b)
|
|
debug_estim_b.append(estim_b)
|
|
|
|
if breaking:
|
|
break
|
|
|
|
# Plot in case we did not reach convergence
|
|
if not breaking:
|
|
import matplotlib.pyplot as plt
|
|
|
|
print(
|
|
"ERROR: It seems that the calibration have not reached convergence. Here are some plot to understand why:")
|
|
print("If you notice unstability, reduce the expected_N value")
|
|
print("If convergece is too slow, increase the expected_N value")
|
|
|
|
plt.figure()
|
|
plt.plot(debug_in)
|
|
plt.plot(debug_out)
|
|
|
|
plt.figure()
|
|
plt.plot(debug_b)
|
|
plt.plot(debug_estim_b)
|
|
|
|
plt.show()
|
|
|
|
a = 1 / 0
|
|
|
|
# Use collected neighbor histogram to get neighbors limit
|
|
cumsum = np.cumsum(neighb_hists.T, axis=0)
|
|
percentiles = np.sum(cumsum < (untouched_ratio * cumsum[hist_n - 1, :]), axis=0)
|
|
self.dataset.neighborhood_limits = percentiles
|
|
|
|
if verbose:
|
|
|
|
# Crop histogram
|
|
while np.sum(neighb_hists[:, -1]) == 0:
|
|
neighb_hists = neighb_hists[:, :-1]
|
|
hist_n = neighb_hists.shape[1]
|
|
|
|
print('\n**************************************************\n')
|
|
line0 = 'neighbors_num '
|
|
for layer in range(neighb_hists.shape[0]):
|
|
line0 += '| layer {:2d} '.format(layer)
|
|
print(line0)
|
|
for neighb_size in range(hist_n):
|
|
line0 = ' {:4d} '.format(neighb_size)
|
|
for layer in range(neighb_hists.shape[0]):
|
|
if neighb_size > percentiles[layer]:
|
|
color = bcolors.FAIL
|
|
else:
|
|
color = bcolors.OKGREEN
|
|
line0 += '|{:}{:10d}{:} '.format(color,
|
|
neighb_hists[layer, neighb_size],
|
|
bcolors.ENDC)
|
|
|
|
print(line0)
|
|
|
|
print('\n**************************************************\n')
|
|
print('\nchosen neighbors limits: ', percentiles)
|
|
print()
|
|
|
|
# Save batch_limit dictionary
|
|
if self.dataset.use_potentials:
|
|
sampler_method = 'potentials'
|
|
else:
|
|
sampler_method = 'random'
|
|
key = '{:s}_{:.3f}_{:.3f}_{:d}'.format(sampler_method,
|
|
self.dataset.config.in_radius,
|
|
self.dataset.config.first_subsampling_dl,
|
|
self.dataset.config.batch_num)
|
|
batch_lim_dict[key] = float(self.dataset.batch_limit)
|
|
with open(batch_lim_file, 'wb') as file:
|
|
pickle.dump(batch_lim_dict, file)
|
|
|
|
# Save neighb_limit dictionary
|
|
for layer_ind in range(self.dataset.config.num_layers):
|
|
dl = self.dataset.config.first_subsampling_dl * (2 ** layer_ind)
|
|
if self.dataset.config.deform_layers[layer_ind]:
|
|
r = dl * self.dataset.config.deform_radius
|
|
else:
|
|
r = dl * self.dataset.config.conv_radius
|
|
key = '{:.3f}_{:.3f}'.format(dl, r)
|
|
neighb_lim_dict[key] = self.dataset.neighborhood_limits[layer_ind]
|
|
with open(neighb_lim_file, 'wb') as file:
|
|
pickle.dump(neighb_lim_dict, file)
|
|
|
|
print('Calibration done in {:.1f}s\n'.format(time.time() - t0))
|
|
return
|
|
|
|
|
|
class NPM3DCustomBatch:
|
|
"""Custom batch definition with memory pinning for NPM3D"""
|
|
|
|
def __init__(self, input_list):
|
|
|
|
# Get rid of batch dimension
|
|
input_list = input_list[0]
|
|
|
|
# Number of layers
|
|
L = (len(input_list) - 7) // 5
|
|
|
|
# Extract input tensors from the list of numpy array
|
|
ind = 0
|
|
self.points = [torch.from_numpy(nparray) for nparray in input_list[ind:ind + L]]
|
|
ind += L
|
|
self.neighbors = [torch.from_numpy(nparray) for nparray in input_list[ind:ind + L]]
|
|
ind += L
|
|
self.pools = [torch.from_numpy(nparray) for nparray in input_list[ind:ind + L]]
|
|
ind += L
|
|
self.upsamples = [torch.from_numpy(nparray) for nparray in input_list[ind:ind + L]]
|
|
ind += L
|
|
self.lengths = [torch.from_numpy(nparray) for nparray in input_list[ind:ind + L]]
|
|
ind += L
|
|
self.features = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.labels = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.scales = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.rots = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.cloud_inds = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.center_inds = torch.from_numpy(input_list[ind])
|
|
ind += 1
|
|
self.input_inds = torch.from_numpy(input_list[ind])
|
|
|
|
return
|
|
|
|
def pin_memory(self):
|
|
"""
|
|
Manual pinning of the memory
|
|
"""
|
|
|
|
self.points = [in_tensor.pin_memory() for in_tensor in self.points]
|
|
self.neighbors = [in_tensor.pin_memory() for in_tensor in self.neighbors]
|
|
self.pools = [in_tensor.pin_memory() for in_tensor in self.pools]
|
|
self.upsamples = [in_tensor.pin_memory() for in_tensor in self.upsamples]
|
|
self.lengths = [in_tensor.pin_memory() for in_tensor in self.lengths]
|
|
self.features = self.features.pin_memory()
|
|
self.labels = self.labels.pin_memory()
|
|
self.scales = self.scales.pin_memory()
|
|
self.rots = self.rots.pin_memory()
|
|
self.cloud_inds = self.cloud_inds.pin_memory()
|
|
self.center_inds = self.center_inds.pin_memory()
|
|
self.input_inds = self.input_inds.pin_memory()
|
|
|
|
return self
|
|
|
|
def to(self, device):
|
|
|
|
self.points = [in_tensor.to(device) for in_tensor in self.points]
|
|
self.neighbors = [in_tensor.to(device) for in_tensor in self.neighbors]
|
|
self.pools = [in_tensor.to(device) for in_tensor in self.pools]
|
|
self.upsamples = [in_tensor.to(device) for in_tensor in self.upsamples]
|
|
self.lengths = [in_tensor.to(device) for in_tensor in self.lengths]
|
|
self.features = self.features.to(device)
|
|
self.labels = self.labels.to(device)
|
|
self.scales = self.scales.to(device)
|
|
self.rots = self.rots.to(device)
|
|
self.cloud_inds = self.cloud_inds.to(device)
|
|
self.center_inds = self.center_inds.to(device)
|
|
self.input_inds = self.input_inds.to(device)
|
|
|
|
return self
|
|
|
|
def unstack_points(self, layer=None):
|
|
"""Unstack the points"""
|
|
return self.unstack_elements('points', layer)
|
|
|
|
def unstack_neighbors(self, layer=None):
|
|
"""Unstack the neighbors indices"""
|
|
return self.unstack_elements('neighbors', layer)
|
|
|
|
def unstack_pools(self, layer=None):
|
|
"""Unstack the pooling indices"""
|
|
return self.unstack_elements('pools', layer)
|
|
|
|
def unstack_elements(self, element_name, layer=None, to_numpy=True):
|
|
"""
|
|
Return a list of the stacked elements in the batch at a certain layer. If no layer is given, then return all
|
|
layers
|
|
"""
|
|
|
|
if element_name == 'points':
|
|
elements = self.points
|
|
elif element_name == 'neighbors':
|
|
elements = self.neighbors
|
|
elif element_name == 'pools':
|
|
elements = self.pools[:-1]
|
|
else:
|
|
raise ValueError('Unknown element name: {:s}'.format(element_name))
|
|
|
|
all_p_list = []
|
|
for layer_i, layer_elems in enumerate(elements):
|
|
|
|
if layer is None or layer == layer_i:
|
|
|
|
i0 = 0
|
|
p_list = []
|
|
if element_name == 'pools':
|
|
lengths = self.lengths[layer_i + 1]
|
|
else:
|
|
lengths = self.lengths[layer_i]
|
|
|
|
for b_i, length in enumerate(lengths):
|
|
|
|
elem = layer_elems[i0:i0 + length]
|
|
if element_name == 'neighbors':
|
|
elem[elem >= self.points[layer_i].shape[0]] = -1
|
|
elem[elem >= 0] -= i0
|
|
elif element_name == 'pools':
|
|
elem[elem >= self.points[layer_i].shape[0]] = -1
|
|
elem[elem >= 0] -= torch.sum(self.lengths[layer_i][:b_i])
|
|
i0 += length
|
|
|
|
if to_numpy:
|
|
p_list.append(elem.numpy())
|
|
else:
|
|
p_list.append(elem)
|
|
|
|
if layer == layer_i:
|
|
return p_list
|
|
|
|
all_p_list.append(p_list)
|
|
|
|
return all_p_list
|
|
|
|
|
|
def NPM3DCollate(batch_data):
|
|
return NPM3DCustomBatch(batch_data)
|
|
|
|
|
|
# ----------------------------------------------------------------------------------------------------------------------
|
|
#
|
|
# Debug functions
|
|
# \*********************/
|
|
|
|
|
|
def debug_upsampling(dataset, loader):
|
|
"""Shows which labels are sampled according to strategy chosen"""
|
|
|
|
for epoch in range(10):
|
|
|
|
for batch_i, batch in enumerate(loader):
|
|
pc1 = batch.points[1].numpy()
|
|
pc2 = batch.points[2].numpy()
|
|
up1 = batch.upsamples[1].numpy()
|
|
|
|
print(pc1.shape, '=>', pc2.shape)
|
|
print(up1.shape, np.max(up1))
|
|
|
|
pc2 = np.vstack((pc2, np.zeros_like(pc2[:1, :])))
|
|
|
|
# Get neighbors distance
|
|
p0 = pc1[10, :]
|
|
neighbs0 = up1[10, :]
|
|
neighbs0 = pc2[neighbs0, :] - p0
|
|
d2 = np.sum(neighbs0 ** 2, axis=1)
|
|
|
|
print(neighbs0.shape)
|
|
print(neighbs0[:5])
|
|
print(d2[:5])
|
|
|
|
print('******************')
|
|
print('*******************************************')
|
|
|
|
_, counts = np.unique(dataset.input_labels, return_counts=True)
|
|
print(counts)
|
|
|
|
|
|
def debug_timing(dataset, loader):
|
|
"""Timing of generator function"""
|
|
|
|
t = [time.time()]
|
|
last_display = time.time()
|
|
mean_dt = np.zeros(2)
|
|
estim_b = dataset.config.batch_num
|
|
estim_N = 0
|
|
|
|
for epoch in range(10):
|
|
|
|
for batch_i, batch in enumerate(loader):
|
|
# print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes)
|
|
|
|
# New time
|
|
t = t[-1:]
|
|
t += [time.time()]
|
|
|
|
# Update estim_b (low pass filter)
|
|
estim_b += (len(batch.cloud_inds) - estim_b) / 100
|
|
estim_N += (batch.features.shape[0] - estim_N) / 10
|
|
|
|
# Pause simulating computations
|
|
time.sleep(0.05)
|
|
t += [time.time()]
|
|
|
|
# Average timing
|
|
mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1]))
|
|
|
|
# Console display (only one per second)
|
|
if (t[-1] - last_display) > -1.0:
|
|
last_display = t[-1]
|
|
message = 'Step {:08d} -> (ms/batch) {:8.2f} {:8.2f} / batch = {:.2f} - {:.0f}'
|
|
print(message.format(batch_i,
|
|
1000 * mean_dt[0],
|
|
1000 * mean_dt[1],
|
|
estim_b,
|
|
estim_N))
|
|
|
|
print('************* Epoch ended *************')
|
|
|
|
_, counts = np.unique(dataset.input_labels, return_counts=True)
|
|
print(counts)
|
|
|
|
|
|
def debug_show_clouds(dataset, loader):
|
|
for epoch in range(10):
|
|
|
|
clouds = []
|
|
cloud_normals = []
|
|
cloud_labels = []
|
|
|
|
L = dataset.config.num_layers
|
|
|
|
for batch_i, batch in enumerate(loader):
|
|
|
|
# Print characteristics of input tensors
|
|
print('\nPoints tensors')
|
|
for i in range(L):
|
|
print(batch.points[i].dtype, batch.points[i].shape)
|
|
print('\nNeigbors tensors')
|
|
for i in range(L):
|
|
print(batch.neighbors[i].dtype, batch.neighbors[i].shape)
|
|
print('\nPools tensors')
|
|
for i in range(L):
|
|
print(batch.pools[i].dtype, batch.pools[i].shape)
|
|
print('\nStack lengths')
|
|
for i in range(L):
|
|
print(batch.lengths[i].dtype, batch.lengths[i].shape)
|
|
print('\nFeatures')
|
|
print(batch.features.dtype, batch.features.shape)
|
|
print('\nLabels')
|
|
print(batch.labels.dtype, batch.labels.shape)
|
|
print('\nAugment Scales')
|
|
print(batch.scales.dtype, batch.scales.shape)
|
|
print('\nAugment Rotations')
|
|
print(batch.rots.dtype, batch.rots.shape)
|
|
print('\nModel indices')
|
|
print(batch.model_inds.dtype, batch.model_inds.shape)
|
|
|
|
print('\nAre input tensors pinned')
|
|
print(batch.neighbors[0].is_pinned())
|
|
print(batch.neighbors[-1].is_pinned())
|
|
print(batch.points[0].is_pinned())
|
|
print(batch.points[-1].is_pinned())
|
|
print(batch.labels.is_pinned())
|
|
print(batch.scales.is_pinned())
|
|
print(batch.rots.is_pinned())
|
|
print(batch.model_inds.is_pinned())
|
|
|
|
show_input_batch(batch)
|
|
|
|
print('*******************************************')
|
|
|
|
_, counts = np.unique(dataset.input_labels, return_counts=True)
|
|
print(counts)
|
|
|
|
|
|
def debug_batch_and_neighbors_calib(dataset, loader):
|
|
"""Timing of generator function"""
|
|
|
|
t = [time.time()]
|
|
last_display = time.time()
|
|
mean_dt = np.zeros(2)
|
|
|
|
for epoch in range(10):
|
|
|
|
for batch_i, input_list in enumerate(loader):
|
|
# print(batch_i, tuple(points.shape), tuple(normals.shape), labels, indices, in_sizes)
|
|
|
|
# New time
|
|
t = t[-1:]
|
|
t += [time.time()]
|
|
|
|
# Pause simulating computations
|
|
time.sleep(0.01)
|
|
t += [time.time()]
|
|
|
|
# Average timing
|
|
mean_dt = 0.9 * mean_dt + 0.1 * (np.array(t[1:]) - np.array(t[:-1]))
|
|
|
|
# Console display (only one per second)
|
|
if (t[-1] - last_display) > 1.0:
|
|
last_display = t[-1]
|
|
message = 'Step {:08d} -> Average timings (ms/batch) {:8.2f} {:8.2f} '
|
|
print(message.format(batch_i,
|
|
1000 * mean_dt[0],
|
|
1000 * mean_dt[1]))
|
|
|
|
print('************* Epoch ended *************')
|
|
|
|
_, counts = np.unique(dataset.input_labels, return_counts=True)
|
|
print(counts)
|