def get_sung_info(data_path, house_names0): data_path = os.path.join(data_path, 'houses') house_names1 = os.listdir(data_path) house_names = [h for h in house_names1 if h in house_names0] infos = [] for house in house_names: house_path = os.path.join(data_path, house) object_path = os.path.join(house_path, 'objects') pcl_fns = glob.glob(os.path.join(house_path, 'pcl*.bin')) pcl_fns.sort() #pcl_fns = [pcl_fn.split('houses')[1] for pcl_fn in pcl_fns] pcl_num = len(pcl_fns) box_fns = glob.glob(os.path.join(object_path, '*.bin')) box_fns.sort() objects = set([os.path.basename(fn).split('_')[0] for fn in box_fns]) for i in range(pcl_num): info = {} info['velodyne_path'] = pcl_fns[i] info['pointcloud_num_features'] = 6 info['image_idx'] = '0' info['image_path'] = 'empty' info['calib/R0_rect'] = np.eye(4) info['calib/Tr_velo_to_cam'] = np.eye(4) # np.array([[0,-1,0,0]. [0,0,-1,0], [1,0,0,0], [0,0,0,1]], dtype=np.float32) info['calib/P2'] = np.eye(4) base_name = os.path.splitext( os.path.basename(pcl_fns[i]) )[0] idx = int(base_name.split('_')[-1]) annos = defaultdict(list) for obj in objects: box_fn = os.path.join(object_path, obj+'_'+str(idx)+'.bin') box = np.fromfile(box_fn, np.float32) box = box.reshape([-1,7]) box = Bbox3D.convert_to_yx_zb_boxes(box) box_num = box.shape[0] annos['location'].append(box[:,0:3]) annos['dimensions'].append(box[:,3:6]) annos['rotation_y'].append(box[:,6]) annos['name'].append( np.array([obj]*box_num) ) annos['difficulty'].append(np.array(['A']*box_num)) annos['bbox'].append( box3d_t_2d(box, info['calib/P2'] ) ) annos['box3d_camera'].append( get_box3d_cam(box, info['calib/R0_rect'], info['calib/Tr_velo_to_cam']) ) bn = box.shape[0] annos["truncated"].append(np.array([0.0]*bn)) annos["occluded"].append(np.array([0.0]*bn)) annos["alpha"].append( get_alpha(box) ) for key in annos: annos[key] = np.concatenate(annos[key], 0) info['annos'] = annos infos.append(info) return infos
def __getitem__(self, index): is_train = self.is_train scale = self.scale full_scale = self.full_scale objects_to_detect = self.objects_to_detect zoom_rate = 0.1 * 0 flip_x = False and is_train random_rotate = False and is_train distortion = False and is_train origin_offset = False and is_train norm_noise = 0.01 * int(is_train) * 0 fn = self.files[index] print(f'Loading {fn} in SUNCGDataset') hn = os.path.basename(os.path.dirname(fn)) #if self.is_train: # print(f'\n(suncg_dataset.py) train {index}-th {hn}\n') #else: # print(f'\n(suncg_dataset.py) test {index}-th {hn}\n') pcl_i, bboxes_dic_i_0 = torch.load(fn) #points_sample(pcl_i) for obj in bboxes_dic_i_0: if type(bboxes_dic_i_0[obj]) == torch.Tensor: bboxes_dic_i_0[obj] = bboxes_dic_i_0[obj].data.numpy() a = pcl_i[:, 0:3].copy() b = pcl_i bboxes_dic_i = {} for obj in objects_to_detect: if not (obj in bboxes_dic_i_0 or obj == 'background'): print(f"unknow class {obj}") import pdb pdb.set_trace() # XXX BREAKPOINT assert False for obj in bboxes_dic_i_0: if ('all' in objects_to_detect) or (obj in objects_to_detect): bboxes_dic_i[obj] = Bbox3D.convert_to_yx_zb_boxes( bboxes_dic_i_0[obj]) if obj in ['ceiling', 'floor', 'room']: bboxes_dic_i[obj] = Bbox3D.set_yaw_zero(bboxes_dic_i[obj]) if SHOW_RAW_INPUT: show_pcl_boxdic(pcl_i, bboxes_dic_i) #--------------------------------------------------------------------- # augmentation of xyz m = np.eye(3) + np.random.randn(3, 3) * zoom_rate # aug: zoom if flip_x: m[0][0] *= np.random.randint(0, 2) * 2 - 1 # aug: x flip m *= scale if random_rotate: theta = np.random.rand() * 2 * math.pi # rotation aug m = np.matmul( m, [[math.cos(theta), math.sin(theta), 0], [-math.sin(theta), math.cos(theta), 0], [0, 0, 1]]) a = np.matmul(a, m) if distortion: a = elastic(a, 6 * scale // 50, 40 * scale / 50) a = elastic(a, 20 * scale // 50, 160 * scale / 50) m = a.min(0) M = a.max(0) q = M - m # aug: the centroid between [0,full_scale] offset = -m if origin_offset: offset += np.clip(full_scale - M + m - 0.001, 0, None) * np.random.rand(3) + np.clip( full_scale - M + m + 0.001, None, 0) * np.random.rand(3) a += offset xyz_min = a.min(0) / scale xyz_max = a.max(0) / scale size3d = np.expand_dims(np.concatenate([xyz_min, xyz_max], 0), 0).astype(np.float32) size3d = torch.from_numpy(size3d) #--------------------------------------------------------------------- # augmentation of feature # aug norm b[:, 3:6] += np.random.randn(3) * norm_noise #--------------------------------------------------------------------- # get elements b = b[:, self.elements_ids] if 'xyz' in self.elements: # import augmentation of xyz to feature b[:, 0:3] = a / scale #--------------------------------------------------------------------- # augment gt boxes for obj in bboxes_dic_i: #print(bboxes_dic_i[obj][:,0:3]) bboxes_dic_i[obj][:, 0:3] += np.expand_dims(offset, 0) / scale #print(bboxes_dic_i[obj][:,0:3]) pass #--------------------------------------------------------------------- assert a.min() >= 0, f"point location should not < 0: {a.min()}" up_check = np.all(a < full_scale[np.newaxis, :], 1) if not np.all(up_check): max_scale = a.max(0) print(f'file: {self.files[index]}') print( f'\nmax scale: {max_scale} > full_scale: {full_scale}, some points will be missed\n' ) if not ENABLE_POINTS_MISSED: import pdb pdb.set_trace() # XXX BREAKPOINT assert False idxs = (a.min(1) >= 0) * (up_check) a = a[idxs] b = b[idxs] #c=c[idxs] a = torch.from_numpy(a).long() #locs = torch.cat([a,torch.LongTensor(a.shape[0],1).fill_(index)],1) locs = a feats = torch.from_numpy(b) #--------------------------------------------------------------------- bboxlist3d = bbox_dic_to_BoxList3D(bboxes_dic_i, size3d, self.dset_metas) labels = bboxlist3d if SHOW_AUG_INPUT: show_pcl_boxdic(pcl_i, bboxes_dic_i) bboxlist3d.show() import pdb pdb.set_trace() # XXX BREAKPOINT pass #batch_scopes(locs, scale) data = {'x': [locs, feats], 'y': labels, 'id': index, 'fn': fn} return data