Commit 6f22a6d3 authored by liyinhao's avatar liyinhao
Browse files

add loading, change names, mmcv, empty line

parent 5219773a
import os.path as osp
import mmcv
import numpy as np
from mmdet.datasets.registry import PIPELINES
@PIPELINES.register_module
class LoadPointsFromFile(object):
def __init__(self, points_dim=4, with_reflectivity=True):
self.points_dim = points_dim
self.with_reflectivity = with_reflectivity
def __call__(self, results):
if results['pts_prefix'] is not None:
filename = osp.join(results['pts_prefix'],
results['img_info']['filename'])
else:
filename = results['img_info']['filename']
points = np.fromfile(
filename, dtype=np.float32).reshape(-1, self.points_dim)
results['points'] = points
return results
def __repr__(self):
repr_str = self.__class__.__name__
repr_str += '(points_dim={})'.format(self.points_dim)
repr_str += '(points_dim={})'.format(self.with_reflectivity)
return repr_str
@PIPELINES.register_module
class LoadMultiViewImageFromFiles(object):
""" Load multi channel images from a list of separate channel files.
Expects results['filename'] to be a list of filenames
"""
def __init__(self, to_float32=False, color_type='unchanged'):
self.to_float32 = to_float32
self.color_type = color_type
def __call__(self, results):
if results['img_prefix'] is not None:
filename = [
osp.join(results['img_prefix'], fname)
for fname in results['img_info']['filename']
]
else:
filename = results['img_info']['filename']
img = np.stack(
[mmcv.imread(name, self.color_type) for name in filename], axis=-1)
if self.to_float32:
img = img.astype(np.float32)
results['filename'] = filename
results['img'] = img
results['img_shape'] = img.shape
results['ori_shape'] = img.shape
# Set initial values for default meta_keys
results['pad_shape'] = img.shape
results['scale_factor'] = 1.0
num_channels = 1 if len(img.shape) < 3 else img.shape[2]
results['img_norm_cfg'] = dict(
mean=np.zeros(num_channels, dtype=np.float32),
std=np.ones(num_channels, dtype=np.float32),
to_rgb=False)
return results
def __repr__(self):
return "{} (to_float32={}, color_type='{}')".format(
self.__class__.__name__, self.to_float32, self.color_type)
...@@ -21,10 +21,10 @@ def create_indoor_info_file(data_path, ...@@ -21,10 +21,10 @@ def create_indoor_info_file(data_path,
""" """
assert os.path.exists(data_path) assert os.path.exists(data_path)
assert pkl_prefix in ['sunrgbd', 'scannet'] assert pkl_prefix in ['sunrgbd', 'scannet']
if save_path is None: save_path = data_path if save_path is None else save_path
save_path = data_path
assert os.path.exists(save_path) assert os.path.exists(save_path)
train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl') train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl')
val_filename = os.path.join(save_path, f'{pkl_prefix}_infos_val.pkl') val_filename = os.path.join(save_path, f'{pkl_prefix}_infos_val.pkl')
if pkl_prefix == 'sunrgbd': if pkl_prefix == 'sunrgbd':
train_dataset = SUNRGBDData( train_dataset = SUNRGBDData(
...@@ -35,10 +35,21 @@ def create_indoor_info_file(data_path, ...@@ -35,10 +35,21 @@ def create_indoor_info_file(data_path,
train_dataset = ScanNetData(root_path=data_path, split='train') train_dataset = ScanNetData(root_path=data_path, split='train')
val_dataset = ScanNetData(root_path=data_path, split='val') val_dataset = ScanNetData(root_path=data_path, split='val')
infos_train = train_dataset.get_infos(has_label=True) infos_train = train_dataset.get_infos(has_label=True)
with open(train_filename, 'wb') as f:
mmcv.dump(infos_train, f, 'pkl') mmcv.dump(infos_train, train_filename, 'pkl')
print(f'{pkl_prefix} info train file is saved to {train_filename}') print(f'{pkl_prefix} info train file is saved to {train_filename}')
infos_val = val_dataset.get_infos(has_label=True) infos_val = val_dataset.get_infos(has_label=True)
with open(val_filename, 'wb') as f:
mmcv.dump(infos_val, f, 'pkl') mmcv.dump(infos_val, val_filename, 'pkl')
print(f'{pkl_prefix} info val file is saved to {val_filename}') print(f'{pkl_prefix} info val file is saved to {val_filename}')
if __name__ == '__main__':
create_indoor_info_file(
data_path='./data/scannet',
pkl_prefix='scannet',
save_path='./data/scannet')
create_indoor_info_file(
data_path='./data/sunrgbd/sunrgbd_trainval',
pkl_prefix='sunrgbd',
save_path='./data/sunrgbd')
...@@ -6,28 +6,28 @@ import numpy as np ...@@ -6,28 +6,28 @@ import numpy as np
import scipy.io as sio import scipy.io as sio
def random_sampling(pc, num_points, replace=None, return_choices=False): def random_sampling(points, num_points, replace=None, return_choices=False):
"""Random Sampling. """Random Sampling.
Sampling point cloud to a certain number of points. Sampling point cloud to a certain number of points.
Args: Args:
pc (ndarray): Point cloud. points (ndarray): Point cloud.
num_points (int): The number of samples. num_points (int): The number of samples.
replace (bool): Whether the sample is with or without replacement. replace (bool): Whether the sample is with or without replacement.
return_choices (bool): Whether to return choices. return_choices (bool): Whether to return choices.
Returns: Returns:
pc (ndarray): Point cloud after sampling. points (ndarray): Point cloud after sampling.
""" """
if replace is None: if replace is None:
replace = (pc.shape[0] < num_points) replace = (points.shape[0] < num_points)
choices = np.random.choice(pc.shape[0], num_points, replace=replace) choices = np.random.choice(points.shape[0], num_points, replace=replace)
if return_choices: if return_choices:
return pc[choices], choices return points[choices], choices
else: else:
return pc[choices] return points[choices]
class SUNRGBDInstance(object): class SUNRGBDInstance(object):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment