thirdparty/learning3d/data_utils/dataloaders.py

import torch
import torch.nn as nn
import torch.nn.functional as F 
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import numpy as np
import os
import h5py
import subprocess
import shlex
import json
import glob
from .. ops import transform_functions, se3
from sklearn.neighbors import NearestNeighbors
from scipy.spatial.distance import minkowski
from scipy.spatial import cKDTree
from torch.utils.data import Dataset

def download_modelnet40():
	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
	if not os.path.exists(DATA_DIR):
		os.mkdir(DATA_DIR)
	if not os.path.exists(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048')):
		www = 'https://shapenet.cs.stanford.edu/media/modelnet40_ply_hdf5_2048.zip'
		zipfile = os.path.basename(www)
		os.system('wget --no-check-certificate %s; unzip %s' % (www, zipfile))
		os.system('mv %s %s' % (zipfile[:-4], DATA_DIR))
		os.system('rm %s' % (zipfile))

def load_data(train, use_normals):
	if train: partition = 'train'
	else: partition = 'test'
	BASE_DIR = os.path.dirname(os.path.abspath(__file__))
	DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
	all_data = []
	all_label = []
	for h5_name in glob.glob(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'ply_data_%s*.h5' % partition)):
		f = h5py.File(h5_name)
		if use_normals: data = np.concatenate([f['data'][:], f['normal'][:]], axis=-1).astype('float32')
		else: data = f['data'][:].astype('float32')
		label = f['label'][:].astype('int64')
		f.close()
		all_data.append(data)
		all_label.append(label)
	all_data = np.concatenate(all_data, axis=0)
	all_label = np.concatenate(all_label, axis=0)
	return all_data, all_label

def deg_to_rad(deg):
	return np.pi / 180 * deg

def create_random_transform(dtype, max_rotation_deg, max_translation):
	max_rotation = deg_to_rad(max_rotation_deg)
	rot = np.random.uniform(-max_rotation, max_rotation, [1, 3])
	trans = np.random.uniform(-max_translation, max_translation, [1, 3])
	quat = transform_functions.euler_to_quaternion(rot, "xyz")

	vec = np.concatenate([quat, trans], axis=1)
	vec = torch.tensor(vec, dtype=dtype)
	return vec

def jitter_pointcloud(pointcloud, sigma=0.04, clip=0.05):
	# N, C = pointcloud.shape
	sigma = 0.04*np.random.random_sample()
	pointcloud += torch.empty(pointcloud.shape).normal_(mean=0, std=sigma).clamp(-clip, clip)
	return pointcloud

def farthest_subsample_points(pointcloud1, num_subsampled_points=768):
	pointcloud1 = pointcloud1
	num_points = pointcloud1.shape[0]
	nbrs1 = NearestNeighbors(n_neighbors=num_subsampled_points, algorithm='auto',
							 metric=lambda x, y: minkowski(x, y)).fit(pointcloud1[:, :3])
	random_p1 = np.random.random(size=(1, 3)) + np.array([[500, 500, 500]]) * np.random.choice([1, -1, 1, -1])
	idx1 = nbrs1.kneighbors(random_p1, return_distance=False).reshape((num_subsampled_points,))
	gt_mask = torch.zeros(num_points).scatter_(0, torch.tensor(idx1), 1)
	return pointcloud1[idx1, :], gt_mask

def uniform_2_sphere(num: int = None):
	"""Uniform sampling on a 2-sphere

	Source: https://gist.github.com/andrewbolster/10274979

	Args:
		num: Number of vectors to sample (or None if single)

	Returns:
		Random Vector (np.ndarray) of size (num, 3) with norm 1.
		If num is None returned value will have size (3,)

	"""
	if num is not None:
		phi = np.random.uniform(0.0, 2 * np.pi, num)
		cos_theta = np.random.uniform(-1.0, 1.0, num)
	else:
		phi = np.random.uniform(0.0, 2 * np.pi)
		cos_theta = np.random.uniform(-1.0, 1.0)

	theta = np.arccos(cos_theta)
	x = np.sin(theta) * np.cos(phi)
	y = np.sin(theta) * np.sin(phi)
	z = np.cos(theta)

	return np.stack((x, y, z), axis=-1)

def planar_crop(points, p_keep= 0.7):
	p_keep = np.array(p_keep, dtype=np.float32)

	rand_xyz = uniform_2_sphere()
	pts = points.numpy()
	centroid = np.mean(pts[:, :3], axis=0)
	points_centered = pts[:, :3] - centroid

	dist_from_plane = np.dot(points_centered, rand_xyz)

	mask = dist_from_plane > np.percentile(dist_from_plane, (1.0 - p_keep) * 100)
	idx_x = torch.Tensor(np.nonzero(mask))

	return torch.Tensor(pts[mask, :3]), idx_x

def knn_idx(pts, k):
	kdt = cKDTree(pts) 
	_, idx = kdt.query(pts, k=k+1)
	return idx[:, 1:]

def get_rri(pts, k):
	# pts: N x 3, original points
	# q: N x K x 3, nearest neighbors
	q = pts[knn_idx(pts, k)]
	p = np.repeat(pts[:, None], k, axis=1)
	# rp, rq: N x K x 1, norms
	rp = np.linalg.norm(p, axis=-1, keepdims=True)
	rq = np.linalg.norm(q, axis=-1, keepdims=True)
	pn = p / rp
	qn = q / rq
	dot = np.sum(pn * qn, -1, keepdims=True)
	# theta: N x K x 1, angles
	theta = np.arccos(np.clip(dot, -1, 1))
	T_q = q - dot * p
	sin_psi = np.sum(np.cross(T_q[:, None], T_q[:, :, None]) * pn[:, None], -1)
	cos_psi = np.sum(T_q[:, None] * T_q[:, :, None], -1)
	psi = np.arctan2(sin_psi, cos_psi) % (2*np.pi)
	idx = np.argpartition(psi, 1)[:, :, 1:2]
	# phi: N x K x 1, projection angles
	phi = np.take_along_axis(psi, idx, axis=-1)
	feat = np.concatenate([rp, rq, theta, phi], axis=-1)
	return feat.reshape(-1, k * 4)

def get_rri_cuda(pts, k, npts_per_block=1):
	try:
		import pycuda.autoinit
		from pycuda import gpuarray
		from pycuda.compiler import SourceModule
	except Exception as e:
		print("Error raised in pycuda modules! pycuda only works with GPU, ", e)
		raise

	mod_rri = SourceModule(open('rri.cu').read() % (k, npts_per_block))
	rri_cuda = mod_rri.get_function('get_rri_feature')

	N = len(pts)
	pts_gpu = gpuarray.to_gpu(pts.astype(np.float32).ravel())
	k_idx = knn_idx(pts, k)
	k_idx_gpu = gpuarray.to_gpu(k_idx.astype(np.int32).ravel())
	feat_gpu = gpuarray.GPUArray((N * k * 4,), np.float32)

	rri_cuda(pts_gpu, np.int32(N), k_idx_gpu, feat_gpu,
				grid=(((N-1) // npts_per_block)+1, 1),
				block=(npts_per_block, k, 1))

	feat = feat_gpu.get().reshape(N, k * 4).astype(np.float32)
	return feat


class UnknownDataTypeError(Exception):
	def __init__(self, *args):
		if args: self.message = args[0]
		else: self.message = 'Datatype not understood for dataset.'

	def __str__(self):
		return self.message


class ModelNet40Data(Dataset):
	def __init__(
		self,
		train=True,
		num_points=1024,
		download=True,
		randomize_data=False,
		use_normals=False
	):
		super(ModelNet40Data, self).__init__()
		if download: download_modelnet40()
		self.data, self.labels = load_data(train, use_normals)
		if not train: self.shapes = self.read_classes_ModelNet40()
		self.num_points = num_points
		self.randomize_data = randomize_data

	def __getitem__(self, idx):
		if self.randomize_data: current_points = self.randomize(idx)
		else: current_points = self.data[idx].copy()

		current_points = torch.from_numpy(current_points[:self.num_points, :]).float()
		label = torch.from_numpy(self.labels[idx]).type(torch.LongTensor)

		return current_points, label

	def __len__(self):
		return self.data.shape[0]

	def randomize(self, idx):
		pt_idxs = np.arange(0, self.num_points)
		np.random.shuffle(pt_idxs)
		return self.data[idx, pt_idxs].copy()

	def get_shape(self, label):
		return self.shapes[label]

	def read_classes_ModelNet40(self):
		BASE_DIR = os.path.dirname(os.path.abspath(__file__))
		DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
		file = open(os.path.join(DATA_DIR, 'modelnet40_ply_hdf5_2048', 'shape_names.txt'), 'r')
		shape_names = file.read()
		shape_names = np.array(shape_names.split('\n')[:-1])
		return shape_names


class ClassificationData(Dataset):
	def __init__(self, data_class=ModelNet40Data()):
		super(ClassificationData, self).__init__()
		self.set_class(data_class)

	def __len__(self):
		return len(self.data_class)

	def set_class(self, data_class):
		self.data_class = data_class

	def get_shape(self, label):
		try:
			return self.data_class.get_shape(label)
		except:
			return -1

	def __getitem__(self, index):
		return self.data_class[index]


class RegistrationData(Dataset):
	def __init__(self, algorithm, data_class=ModelNet40Data(), partial_source=False, partial_template=False, noise=False, additional_params={}):
		super(RegistrationData, self).__init__()
		available_algorithms = ['PCRNet', 'PointNetLK', 'DCP', 'PRNet', 'iPCRNet', 'RPMNet', 'DeepGMR']
		if algorithm in available_algorithms: self.algorithm = algorithm
		else: raise Exception("Algorithm not available for registration.")
		
		self.set_class(data_class)
		self.partial_template = partial_template
		self.partial_source = partial_source
		self.noise = noise
		self.additional_params = additional_params
		self.use_rri = False

		if self.algorithm == 'PCRNet' or self.algorithm == 'iPCRNet':
			from .. ops.transform_functions import PCRNetTransform
			self.transforms = PCRNetTransform(len(data_class), angle_range=45, translation_range=1)
		if self.algorithm == 'PointNetLK':
			from .. ops.transform_functions import PNLKTransform
			self.transforms = PNLKTransform(0.8, True)
		if self.algorithm == 'RPMNet':
			from .. ops.transform_functions import RPMNetTransform
			self.transforms = RPMNetTransform(0.8, True)
		if self.algorithm == 'DCP' or self.algorithm == 'PRNet':
			from .. ops.transform_functions import DCPTransform
			self.transforms = DCPTransform(angle_range=45, translation_range=1)
		if self.algorithm == 'DeepGMR':
			self.get_rri = get_rri_cuda if torch.cuda.is_available() else get_rri
			from .. ops.transform_functions import DeepGMRTransform
			self.transforms = DeepGMRTransform(angle_range=90, translation_range=1)
			if 'nearest_neighbors' in self.additional_params.keys() and self.additional_params['nearest_neighbors'] > 0:
				self.use_rri = True
				self.nearest_neighbors = self.additional_params['nearest_neighbors']

	def __len__(self):
		return len(self.data_class)

	def set_class(self, data_class):
		self.data_class = data_class

	def __getitem__(self, index):
		template, label = self.data_class[index]
		self.transforms.index = index				# for fixed transformations in PCRNet.
		source = self.transforms(template)

		# Check for Partial Data.
		if self.additional_params.get('partial_point_cloud_method', None) == 'planar_crop':
			source, gt_idx_source = planar_crop(source)
			template, gt_idx_template = planar_crop(template)
			intersect_mask, intersect_x, intersect_y  = np.intersect1d(gt_idx_source, gt_idx_template, return_indices=True)

			self.template_mask = torch.zeros(template.shape[0])
			self.source_mask = torch.zeros(source.shape[0])
			self.template_mask[intersect_y]  = 1
			self.source_mask[intersect_x]  = 1
		else:
			if self.partial_source: source, self.source_mask = farthest_subsample_points(source)
			if self.partial_template: template, self.template_mask = farthest_subsample_points(template)



		# Check for Noise in Source Data.
		if self.noise: source = jitter_pointcloud(source)

		if self.use_rri:
			template, source = template.numpy(), source.numpy()
			template = np.concatenate([template, self.get_rri(template - template.mean(axis=0), self.nearest_neighbors)], axis=1)
			source = np.concatenate([source, self.get_rri(source - source.mean(axis=0), self.nearest_neighbors)], axis=1)
			template, source = torch.tensor(template).float(), torch.tensor(source).float()

		igt = self.transforms.igt
		
		if self.additional_params.get('use_masknet', False):
			if self.partial_source and self.partial_template:
				return template, source, igt, self.template_mask, self.source_mask
			elif self.partial_source:
				return template, source, igt, self.source_mask
			elif self.partial_template:
				return template, source, igt, self.template_mask
		else:
			return template, source, igt


class SegmentationData(Dataset):
	def __init__(self):
		super(SegmentationData, self).__init__()

	def __len__(self):
		pass

	def __getitem__(self, index):
		pass


class FlowData(Dataset):
	def __init__(self):
		super(FlowData, self).__init__()
		self.pc1, self.pc2, self.flow = self.read_data()

	def __len__(self):
		if isinstance(self.pc1, np.ndarray):
			return self.pc1.shape[0]
		elif isinstance(self.pc1, list):
			return len(self.pc1)
		else:
			raise UnknownDataTypeError

	def read_data(self):
		pass

	def __getitem__(self, index):
		return self.pc1[index], self.pc2[index], self.flow[index]


class SceneflowDataset(Dataset):
	def __init__(self, npoints=1024, root='', partition='train'):
		if root == '':
			BASE_DIR = os.path.dirname(os.path.abspath(__file__))
			DATA_DIR = os.path.join(BASE_DIR, os.pardir, 'data')
			root = os.path.join(DATA_DIR, 'data_processed_maxcut_35_20k_2k_8192')
			if not os.path.exists(root): 
				print("To download dataset, click here: https://drive.google.com/file/d/1CMaxdt-Tg1Wct8v8eGNwuT7qRSIyJPY-/view")
				exit()
			else:
				print("SceneflowDataset Found Successfully!")

		self.npoints = npoints
		self.partition = partition
		self.root = root
		if self.partition=='train':
			self.datapath = glob.glob(os.path.join(self.root, 'TRAIN*.npz'))
		else:
			self.datapath = glob.glob(os.path.join(self.root, 'TEST*.npz'))
		self.cache = {}
		self.cache_size = 30000

		###### deal with one bad datapoint with nan value
		self.datapath = [d for d in self.datapath if 'TRAIN_C_0140_left_0006-0' not in d]
		######
		print(self.partition, ': ',len(self.datapath))

	def __getitem__(self, index):
		if index in self.cache:
			pos1, pos2, color1, color2, flow, mask1 = self.cache[index]
		else:
			fn = self.datapath[index]
			with open(fn, 'rb') as fp:
				data = np.load(fp)
				pos1 = data['points1'].astype('float32')
				pos2 = data['points2'].astype('float32')
				color1 = data['color1'].astype('float32')
				color2 = data['color2'].astype('float32')
				flow = data['flow'].astype('float32')
				mask1 = data['valid_mask1']

			if len(self.cache) < self.cache_size:
				self.cache[index] = (pos1, pos2, color1, color2, flow, mask1)

		if self.partition == 'train':
			n1 = pos1.shape[0]
			sample_idx1 = np.random.choice(n1, self.npoints, replace=False)
			n2 = pos2.shape[0]
			sample_idx2 = np.random.choice(n2, self.npoints, replace=False)

			pos1 = pos1[sample_idx1, :]
			pos2 = pos2[sample_idx2, :]
			color1 = color1[sample_idx1, :]
			color2 = color2[sample_idx2, :]
			flow = flow[sample_idx1, :]
			mask1 = mask1[sample_idx1]
		else:
			pos1 = pos1[:self.npoints, :]
			pos2 = pos2[:self.npoints, :]
			color1 = color1[:self.npoints, :]
			color2 = color2[:self.npoints, :]
			flow = flow[:self.npoints, :]
			mask1 = mask1[:self.npoints]

		pos1_center = np.mean(pos1, 0)
		pos1 -= pos1_center
		pos2 -= pos1_center

		return pos1, pos2, color1, color2, flow, mask1

	def __len__(self):
		return len(self.datapath)


if __name__ == '__main__':
	class Data():
		def __init__(self):
			super(Data, self).__init__()
			self.data, self.label = self.read_data()

		def read_data(self):
			return [4,5,6], [4,5,6]

		def __len__(self):
			return len(self.data)

		def __getitem__(self, idx):
			return self.data[idx], self.label[idx]

	cd = RegistrationData('abc')
	import ipdb; ipdb.set_trace()