def __getitem__(self, idx): video = self.file_list[idx] frames = [] masks = [] opt_flows = [] n_objects = [] frame_indexes = self._get_frame_indexes(video['n_frames'], self.n_max_frames) for fi in frame_indexes: frame = np.array(IO.get(video['frames'][fi]).convert('RGB')) frames.append(np.array(frame)) mask = IO.get(video['masks'][fi]) mask = mask.convert('P') if mask is not None else np.zeros( frame.shape[:-1]) masks.append(np.array(mask)) opt_flow = IO.get( video['optical_flow'][fi]) if 'optical_flow' in video else None opt_flow = opt_flow if opt_flow is not None else np.zeros( frame.shape[:-1] + (2, )) opt_flows.append(np.array(opt_flow)) # The number of objects in the masks mask_indexes = set() for m in masks: _mask_indexes = np.unique(m) _mask_indexes = _mask_indexes[_mask_indexes != self.ignore_idx] mask_indexes.update(_mask_indexes) _n_objects = min(len(mask_indexes) - 1, self.n_max_objects) n_objects.append(_n_objects) # Data preprocessing and augmentation if self.transforms is not None: frames, masks, opt_flows = self.transforms(frames, masks, opt_flows) return video['name'], np.array(n_objects), frames, masks, opt_flows
def __getitem__(self, idx): sample = self.file_list[idx] data = {} rand_idx = -1 if 'n_renderings' in self.options: rand_idx = random.randint(0, self.options['n_renderings'] - 1) if self.options['shuffle'] else 0 for ri in self.options['required_items']: file_path = sample['%s_path' % ri] if type(file_path) == list: file_path = file_path[rand_idx] # print(file_path) data[ri] = IO.get(file_path).astype(np.float32) if self.transforms is not None: data = self.transforms(data) return sample['taxonomy_id'], sample['model_id'], data
def __getitem__(self, idx): sample = self.file_list[idx] data = {} # rand_idx = -1 # if 'n_renderings' in self.options: rand_idx = random.randint(0, 8 - 1) if self.options['shuffle'] else 0 # print('rand_idx2: ', rand_idx) for ri in self.options['required_items']: # 'partial' & 'gt' file_path = sample['%s_path' % ri] if type(file_path) == list: file_path = file_path[rand_idx] # print('rand_idx1: ', rand_idx) # 返回值是三维坐标的列表(np) data[ri] = IO.get(file_path).astype(np.float32) # print('data\'s shape: ', data['partial_cloud'].shape) if self.transforms is not None: data = self.transforms(data) return sample['taxonomy_id'], sample['model_id'], data