def __init__(self, dataset_config): """ Initializes directories, and loads the sample list Args: dataset_config: KittiDatasetConfig name: unique name for the dataset data_split: "train", "val", "test", "trainval" data_split_dir: must be specified for custom "training" dataset_dir: Kitti dataset dir if not in default location classes: relevant classes num_clusters: number of k-means clusters to separate for each class """ # Parse config self.config = dataset_config self.name = self.config.name self.data_split = self.config.data_split self.dataset_dir = os.path.expanduser(self.config.dataset_dir) data_split_dir = self.config.data_split_dir self.has_labels = self.config.has_labels self.cluster_split = self.config.cluster_split self.classes = list(self.config.classes) self.num_classes = len(self.classes) self.num_clusters = np.asarray(self.config.num_clusters) self.bev_source = self.config.bev_source self.aug_list = self.config.aug_list # Determines the network mode. This is initialized to 'train' but # is overwritten inside the model based on the mode. self.train_val_test = 'train' # Determines if training includes all samples, including the ones # without anchor_info. This is initialized to False, but is overwritten # via the config inside the model. self.train_on_all_samples = False self._set_up_classes_name() # Check that paths and split are valid self._check_dataset_dir() # Get all files and folders in dataset directory all_files = os.listdir(self.dataset_dir) # Get possible data splits from txt files in dataset folder possible_splits = [] for file_name in all_files: if fnmatch.fnmatch(file_name, '*.txt'): possible_splits.append(os.path.splitext(file_name)[0]) # This directory contains a readme.txt file, remove it from the list if 'readme' in possible_splits: possible_splits.remove('readme') if self.data_split not in possible_splits: raise ValueError( "Invalid data split: {}, possible_splits: {}".format( self.data_split, possible_splits)) # Check data_split_dir # Get possible data split dirs from folder names in dataset folder possible_split_dirs = [] for folder_name in all_files: if os.path.isdir(self.dataset_dir + '/' + folder_name): possible_split_dirs.append(folder_name) if data_split_dir in possible_split_dirs: split_dir = self.dataset_dir + '/' + data_split_dir self._data_split_dir = split_dir else: raise ValueError( "Invalid data split dir: {}, possible dirs".format( data_split_dir, possible_split_dirs)) # Batch pointers self._index_in_epoch = 0 self.epochs_completed = 0 self._cam_idx = 2 # Initialize the sample list loaded_sample_names = self.load_sample_names(self.data_split) # Augment the sample list aug_sample_list = [] # Loop through augmentation lengths e.g. # 0: [] # 1: ['flip'], ['pca_jitter'] # 2: ['flip', 'pca_jitter'] for aug_idx in range(len(self.aug_list) + 1): # Get all combinations augmentations = list(itertools.combinations( self.aug_list, aug_idx)) for augmentation in augmentations: for sample_name in loaded_sample_names: aug_sample_list.append(Sample(sample_name, augmentation)) self.sample_list = np.asarray(aug_sample_list) self.num_samples = len(self.sample_list) self._set_up_directories() # Setup utils object self.kitti_utils = KittiUtils(self)
class KittiDataset: def __init__(self, dataset_config): """ Initializes directories, and loads the sample list Args: dataset_config: KittiDatasetConfig name: unique name for the dataset data_split: "train", "val", "test", "trainval" data_split_dir: must be specified for custom "training" dataset_dir: Kitti dataset dir if not in default location classes: relevant classes num_clusters: number of k-means clusters to separate for each class """ # Parse config self.config = dataset_config self.name = self.config.name self.data_split = self.config.data_split self.dataset_dir = os.path.expanduser(self.config.dataset_dir) data_split_dir = self.config.data_split_dir self.has_labels = self.config.has_labels self.cluster_split = self.config.cluster_split self.classes = list(self.config.classes) self.num_classes = len(self.classes) self.num_clusters = np.asarray(self.config.num_clusters) self.bev_source = self.config.bev_source self.aug_list = self.config.aug_list # Determines the network mode. This is initialized to 'train' but # is overwritten inside the model based on the mode. self.train_val_test = 'train' # Determines if training includes all samples, including the ones # without anchor_info. This is initialized to False, but is overwritten # via the config inside the model. self.train_on_all_samples = False self._set_up_classes_name() # Check that paths and split are valid self._check_dataset_dir() # Get all files and folders in dataset directory all_files = os.listdir(self.dataset_dir) # Get possible data splits from txt files in dataset folder possible_splits = [] for file_name in all_files: if fnmatch.fnmatch(file_name, '*.txt'): possible_splits.append(os.path.splitext(file_name)[0]) # This directory contains a readme.txt file, remove it from the list if 'readme' in possible_splits: possible_splits.remove('readme') if self.data_split not in possible_splits: raise ValueError( "Invalid data split: {}, possible_splits: {}".format( self.data_split, possible_splits)) # Check data_split_dir # Get possible data split dirs from folder names in dataset folder possible_split_dirs = [] for folder_name in all_files: if os.path.isdir(self.dataset_dir + '/' + folder_name): possible_split_dirs.append(folder_name) if data_split_dir in possible_split_dirs: split_dir = self.dataset_dir + '/' + data_split_dir self._data_split_dir = split_dir else: raise ValueError( "Invalid data split dir: {}, possible dirs".format( data_split_dir, possible_split_dirs)) # Batch pointers self._index_in_epoch = 0 self.epochs_completed = 0 self._cam_idx = 2 # Initialize the sample list loaded_sample_names = self.load_sample_names(self.data_split) # Augment the sample list aug_sample_list = [] # Loop through augmentation lengths e.g. # 0: [] # 1: ['flip'], ['pca_jitter'] # 2: ['flip', 'pca_jitter'] for aug_idx in range(len(self.aug_list) + 1): # Get all combinations augmentations = list(itertools.combinations( self.aug_list, aug_idx)) for augmentation in augmentations: for sample_name in loaded_sample_names: aug_sample_list.append(Sample(sample_name, augmentation)) self.sample_list = np.asarray(aug_sample_list) self.num_samples = len(self.sample_list) self._set_up_directories() # Setup utils object self.kitti_utils = KittiUtils(self) # Paths @property def rgb_image_dir(self): return self.image_dir @property def rgb_image_r_dir(self): return self.image_r_dir @property def sample_names(self): # This is a property since the sample list gets shuffled for training return np.asarray([sample.name for sample in self.sample_list]) @property def bev_image_dir(self): raise NotImplementedError("BEV images not saved to disk yet!") def _check_dataset_dir(self): """Checks that dataset directory exists in the file system Raises: FileNotFoundError: if the dataset folder is missing """ # Check that dataset path is valid if not os.path.exists(self.dataset_dir): raise FileNotFoundError('Dataset path does not exist: {}'.format( self.dataset_dir)) def _set_up_directories(self): """Sets up data directories.""" # Setup Directories self.image_dir = self._data_split_dir + '/image_2' self.image_r_dir = self._data_split_dir + '/image_3' self.calib_dir = self._data_split_dir + '/calib' self.disp_dir = self._data_split_dir + '/disparity' self.planes_dir = self._data_split_dir + '/planes' self.velo_dir = self._data_split_dir + '/velodyne' self.depth_dir = self._data_split_dir + '/depth_' + str(self._cam_idx) self.score_2d_dir = self._data_split_dir + '/score' # Labels are always in the training folder self.label_dir = self.dataset_dir + \ '/training/label_' + str(self._cam_idx) def _set_up_classes_name(self): # Unique identifier for multiple classes if self.num_classes > 1: if self.classes == ['Pedestrian', 'Cyclist']: self.classes_name = 'People' elif self.classes == ['Car', 'Pedestrian', 'Cyclist']: self.classes_name = 'All' else: raise NotImplementedError('Need new unique identifier for ' 'multiple classes') else: self.classes_name = self.classes[0] # Get sample paths def get_rgb_image_path(self, sample_name): return self.rgb_image_dir + '/' + sample_name + '.png' def get_rgb_image_r_path(self, sample_name): return self.rgb_image_r_dir + '/' + sample_name + '.png' def get_depth_map_path(self, sample_name): return self.depth_dir + '/' + sample_name + '_left_depth.png' def get_velodyne_path(self, sample_name): return self.velo_dir + '/' + sample_name + '.bin' def get_bev_sample_path(self, sample_name): return self.bev_image_dir + '/' + sample_name + '.png' def get_score_2d_path(self, sample_name): return self.score_2d_dir + '/' + sample_name + '.npy' # Cluster info def get_cluster_info(self): return self.kitti_utils.clusters, self.kitti_utils.std_devs def get_anchors_info(self, sample_name): return self.kitti_utils.get_anchors_info( self.classes_name, self.kitti_utils.anchor_strides, sample_name) # Data loading methods def load_sample_names(self, data_split): """Load the sample names listed in this dataset's set file (e.g. train.txt, validation.txt) Args: data_split: override the sample list to load (e.g. for clustering) Returns: A list of sample names (file names) read from the .txt file corresponding to the data split """ set_file = self.dataset_dir + '/' + data_split + '.txt' with open(set_file, 'r') as f: sample_names = f.read().splitlines() return np.array(sample_names) def load_samples(self, indices): """ Loads input-output data for a set of samples. Should only be called when a particular sample dict is required. Otherwise, samples should be provided by the next_batch function Args: indices: A list of sample indices from the dataset.sample_list to be loaded Return: samples: a list of data sample dicts """ sample_dicts = [] for sample_idx in indices: sample = self.sample_list[sample_idx] sample_name = sample.name # Only read labels if they exist if self.has_labels: # Read mini batch first to see if it is empty anchors_info = self.get_anchors_info(sample_name) label_score_2d_path = self.get_score_2d_path(sample_name) with open(label_score_2d_path, 'rb') as file: label_score_2d = np.load(file) file.close() #label_score_2d = np.zeros((12, 39, 2)) if (not anchors_info) and self.train_val_test == 'train' \ and (not self.train_on_all_samples): empty_sample_dict = { constants.KEY_SAMPLE_NAME: sample_name, constants.KEY_ANCHORS_INFO: anchors_info } return [empty_sample_dict] obj_labels = obj_utils.read_labels(self.label_dir, int(sample_name)) # Only use objects that match dataset classes obj_labels = self.kitti_utils.filter_labels(obj_labels) else: obj_labels = None anchors_info = [] label_anchors = np.zeros((1, 6)) label_boxes_3d = np.zeros((1, 7)) label_classes = np.zeros(1) label_score_2d = np.zeros((12, 39, 2)) img_idx = int(sample_name) # Load image (BGR -> RGB) cv_bgr_image = cv2.imread(self.get_rgb_image_path(sample_name)) rgb_image = cv_bgr_image[..., ::-1] image_shape = rgb_image.shape[0:2] image_input = rgb_image cv_bgr_image_r = cv2.imread(self.get_rgb_image_r_path(sample_name)) rgb_image_r = cv_bgr_image_r[..., ::-1] image_r_input = rgb_image_r image_r_shape = rgb_image_r.shape[0:2] # Get ground plane ground_plane = obj_utils.get_road_plane(int(sample_name), self.planes_dir) # Get calibration stereo_calib_p2 = calib_utils.read_calibration( self.calib_dir, int(sample_name)).p2 stereo_calib_p3 = calib_utils.read_calibration( self.calib_dir, int(sample_name)).p3 """ point_cloud = self.kitti_utils.get_point_cloud(self.bev_source, img_idx, image_shape) """ # Augmentation (Flipping) """ if kitti_aug.AUG_FLIPPING in sample.augs: image_input = kitti_aug.flip_image(image_input) image_r_input = kitti_aug.flip_image(image_r_input) image_input, image_r_input = image_r_input, image_input #ipoint_cloud = kitti_aug.flip_point_cloud(point_cloud) obj_labels = [kitti_aug.flip_label_in_3d_only(obj) for obj in obj_labels] ground_plane = kitti_aug.flip_ground_plane(ground_plane) stereo_calib_p2 = kitti_aug.flip_stereo_calib_p2( stereo_calib_p2, image_shape) stereo_calib_p3 = kitti_aug.flip_stereo_calib_p2( stereo_calib_p3, image_r_shape) stereo_calib_p2, stereo_calib_p3 = stereo_calib_p3, stereo_calib_p2 """ # Augmentation (Image Jitter) if kitti_aug.AUG_PCA_JITTER in sample.augs: #print('Jittering') image_input[:, :, 0:3] = kitti_aug.apply_pca_jitter(image_input[:, :, 0:3]) image_r_input[:, :, 0:3] = kitti_aug.apply_pca_jitter( image_r_input[:, :, 0:3]) if obj_labels is not None: label_boxes_3d = np.asarray([ box_3d_encoder.object_label_to_box_3d(obj_label) for obj_label in obj_labels ]) label_classes = [ self.kitti_utils.class_str_to_index(obj_label.type) for obj_label in obj_labels ] label_classes = np.asarray(label_classes, dtype=np.int32) # Return empty anchors_info if no ground truth after filtering if len(label_boxes_3d) == 0: anchors_info = [] if self.train_on_all_samples: # If training without any positive labels, we cannot # set these to zeros, because later on the offset calc # uses log on these anchors. So setting any arbitrary # number here that does not break the offset calculation # should work, since the negative samples won't be # regressed in any case. dummy_anchors = [[-1000, -1000, -1000, 1, 1, 1]] label_anchors = np.asarray(dummy_anchors) dummy_boxes = [[-1000, -1000, -1000, 1, 1, 1, 0]] label_boxes_3d = np.asarray(dummy_boxes) else: label_anchors = np.zeros((1, 6)) label_boxes_3d = np.zeros((1, 7)) label_classes = np.zeros(1) else: label_anchors = box_3d_encoder.box_3d_to_anchor( label_boxes_3d, ortho_rotate=True) # Create BEV maps """ bev_images = self.kitti_utils.create_bev_maps( point_cloud, ground_plane) height_maps = bev_images.get('height_maps') density_map = bev_images.get('density_map') bev_input = np.dstack((*height_maps, density_map)) """ sample_dict = { constants.KEY_LABEL_BOXES_3D: label_boxes_3d, constants.KEY_LABEL_ANCHORS: label_anchors, constants.KEY_LABEL_CLASSES: label_classes, constants.KEY_LABEL_SCORE_2D: label_score_2d, constants.KEY_IMAGE_INPUT: image_input, constants.KEY_IMAGE_R_INPUT: image_r_input, #constants.KEY_BEV_INPUT: bev_input, constants.KEY_ANCHORS_INFO: anchors_info, #constants.KEY_POINT_CLOUD: point_cloud, constants.KEY_GROUND_PLANE: ground_plane, constants.KEY_STEREO_CALIB_P2: stereo_calib_p2, constants.KEY_STEREO_CALIB_P3: stereo_calib_p3, constants.KEY_SAMPLE_NAME: sample_name, constants.KEY_SAMPLE_AUGS: sample.augs } sample_dicts.append(sample_dict) return sample_dicts def _shuffle_samples(self): perm = np.arange(self.num_samples) np.random.shuffle(perm) self.sample_list = self.sample_list[perm] def next_batch(self, batch_size, shuffle=True): """ Retrieve the next `batch_size` samples from this data set. Args: batch_size: number of samples in the batch shuffle: whether to shuffle the indices after an epoch is completed Returns: list of dictionaries containing sample information """ # Create empty set of samples samples_in_batch = [] start = self._index_in_epoch # Shuffle only for the first epoch if self.epochs_completed == 0 and start == 0 and shuffle: self._shuffle_samples() # Go to the next epoch if start + batch_size >= self.num_samples: # Finished epoch self.epochs_completed += 1 # Get the rest examples in this epoch rest_num_examples = self.num_samples - start # Append those samples to the current batch samples_in_batch.extend( self.load_samples(np.arange(start, self.num_samples))) # Shuffle the data if shuffle: self._shuffle_samples() # Start next epoch start = 0 self._index_in_epoch = batch_size - rest_num_examples end = self._index_in_epoch # Append the rest of the batch samples_in_batch.extend(self.load_samples(np.arange(start, end))) else: self._index_in_epoch += batch_size end = self._index_in_epoch # Append the samples in the range to the batch samples_in_batch.extend(self.load_samples(np.arange(start, end))) return samples_in_batch