def __getitem__(self, idx): file = self.files[idx] name = None if self.is_chunks: name = os.path.splitext(os.path.basename(file))[0] inputs, targets, dims, world2grid, target_known, target_hierarchy = data_util.load_train_file(file) else: input_file = file[0] target_file = file[1] name = os.path.splitext(os.path.basename(input_file))[0] inputs, dims, world2grid = data_util.load_scene(input_file) targets, dims, world2grid = data_util.load_scene(target_file) target_known = data_util.load_scene_known(os.path.splitext(target_file)[0] + '.knw') targets = data_util.sparse_to_dense_np(targets[0], targets[1], dims[2], dims[1], dims[0], -float('inf')) target_hierarchy = None orig_dims = torch.LongTensor(targets.shape) if not self.is_chunks: # add padding hierarchy_factor = pow(2, self.num_hierarchy_levels-1) max_input_dim = np.array(sdf.shape) if self.max_input_height > 0 and max_input_dim[self.UP_AXIS] > self.max_input_height: max_input_dim[self.UP_AXIS] = self.max_input_height mask_input = input[0][:,self.UP_AXIS] < self.max_input_height input[0] = input[0][mask_input] input[1] = input[1][mask_input] max_input_dim = ((max_input_dim + (hierarchy_factor*4) - 1) // (hierarchy_factor*4)) * (hierarchy_factor*4) # pad target to max_input_dim padded = np.zeros((max_input_dim[0], max_input_dim[1], max_input_dim[2]), dtype=np.float32) padded.fill(-float('inf')) padded[:min(self.max_input_height, sdf.shape[0]), :sdf.shape[1], :sdf.shape[2]] = sdf[:self.max_input_height, :, :] sdf = padded if target_known is not None: known_pad = np.ones((max_input_dim[0], max_input_dim[1], max_input_dim[2]), dtype=np.uint8) * 255 known_pad[:min(self.max_input_height,target_known.shape[0]), :target_known.shape[1], :target_known.shape[2]] = target_known[:self.max_input_height, :, :] target_known = known_pad else: if self.num_hierarchy_levels < 4: target_hierarchy = target_hierarchy[4-self.num_hierarchy_levels:] mask = np.abs(inputs[1]) < self.truncation input_locs = inputs[0][mask] input_vals = inputs[1][mask] inputs = [torch.from_numpy(input_locs).long(), torch.from_numpy(input_vals[:,np.newaxis]).float()] targets = targets[np.newaxis,:] targets = torch.from_numpy(targets) if target_hierarchy is not None: for h in range(len(target_hierarchy)): target_hierarchy[h] = torch.from_numpy(target_hierarchy[h][np.newaxis,:]) world2grid = torch.from_numpy(world2grid) target_known = target_known[np.newaxis,:] target_known = torch.from_numpy(target_known) sample = {'name': name, 'input': inputs, 'sdf': targets, 'world2grid': world2grid, 'known': target_known, 'hierarchy': target_hierarchy, 'orig_dims': orig_dims} return sample
def tsdf_callback(self, data): # self.test() # print("hhahahahha") time1 = time.time() gen = point_cloud2.read_points(data.points, field_names=("x", "y", "z", "intensity"), skip_nans=True) center = data.center points = np.array(list(gen)).reshape([-1, 4]) print(points.shape, points[:, -1].max()) offset_x = center[0] - args.voxel_size * args.dimx / 2 offset_y = center[1] - args.voxel_size * args.dimx / 2 offset_z = - 2.7 # points[:, 2].min() -0.1 points[:, 0] -= offset_x points[:, 1] -= offset_y points[:, 2] -= offset_z # print("minz: ", offset_z) points = points / args.voxel_size locs = np.floor(points[:, :3]).astype(np.int32) feats = points[:, 3] # print("feats = ",feats) mask = (locs[:, 0] < args.dimx) & (locs[:, 1] < args.dimx) & (locs[:, 2] < args.dimz) & (locs[:, 0] >= 0) & (locs[:, 1] >= 0) & (locs[:, 2] >= 0) locs = locs[mask] feats = feats[mask] input = data_util.sparse_to_dense_np(locs, feats, args.dimx, args.dimx, args.dimz, -float('inf')) input = input.astype(np.float32) # print("shape after s to d", input.shape) input[input > args.truncation] = args.truncation input[input < -args.truncation] = -args.truncation input.reshape([1, args.dimx, args.dimx, args.dimz]) # input = data_util.tsdf_to_bool(input, trunc=2.0) # for grids whose tsdf value > -2, we consider it as already known known_mask = (input > -2).squeeze() # (80,80,48) original_occ = (np.abs(input) < args.occ_thresord).squeeze() original_free = (input > args.occ_thresord).squeeze() original_unkown = (input < - 2).squeeze() # print("INPUT: %f occ, %f free, %f unkwn"%(original_occ.float().sum()/SUM, original_free.float().sum()/SUM, original_unkown.float().sum()/SUM)) time2 = time.time() print("prepocess time: %fs"%(time2 - time1)) output_occ = self.trt_runner.inference(input) time3 = time.time() print("model time: %fs"%(time3 - time2)) output_occ = np.abs(output_occ) < args.occ_thresord # fix conflicts with input in known grids # embed() output_occ[known_mask] = original_occ[known_mask] # added_occ = output_occ.float() > original_occ.float() added_occ = output_occ > original_occ if output_occ is not None: # added_occ = added_occ.squeeze() # todo: compare of np and torch+cuda occ_coords_1 = POINTER[added_occ].reshape([-1, 3]) num_points = occ_coords_1.shape[0] # added_color_1 = torch.ones(num_points).unsqueeze(1) * 10 added_color_1 = np.ones([num_points, 1], dtype=np.float32) * 10 # occ_coords_1 = torch.cat((occ_coords_1, added_color_1), 1) occ_coords_1 = np.concatenate((occ_coords_1, added_color_1), 1) # occ_coords = torch.cat((occ_coords, occ_coords_1), 0) occ_coords = occ_coords_1 print("added size: ",occ_coords.shape, added_occ.shape) if occ_coords.shape[0] > 0: occ_coords *= args.voxel_size occ_coords[:, 0] += offset_x occ_coords[:, 1] += offset_y occ_coords[:, 2] += offset_z # occ_coords = occ_coords.cpu().numpy() msg = data.points msg.height = 1 msg.width = occ_coords.shape[0] msg.fields = [ PointField('x', 0, PointField.FLOAT32, 1), PointField('y', 4, PointField.FLOAT32, 1), PointField('z', 8, PointField.FLOAT32, 1), PointField('intensity', 12, PointField.FLOAT32, 1)] msg.is_bigendian = False msg.point_step = 16 #12 msg.row_step = 16 * occ_coords.shape[0] msg.is_dense = int(np.isfinite(occ_coords).all()) msg.data = occ_coords.tostring() #np.asarray(occ_coords, np.float32).tostring() print("msg: ", occ_coords.shape, len(msg.data), msg.row_step) # print("postpocess time: %fs"%(time.time() - time3)) self.added_occ_pub.publish(msg) #publish added occ # original_unkown = original_input <= -2 # original_unkown = original_unkown.squeeze() # original_occ = original_occ.squeeze() occ_coords = POINTER[original_occ].reshape([-1, 3]) print("known occ size: ",occ_coords.shape) if occ_coords.shape[0] > 0: occ_coords *= args.voxel_size occ_coords[:, 0] += offset_x occ_coords[:, 1] += offset_y occ_coords[:, 2] += offset_z # occ_coords = occ_coords.cpu().numpy() msg2 = data.points msg2.height = 1 msg2.width = occ_coords.shape[0] msg2.fields = [ PointField('x', 0, PointField.FLOAT32, 1), PointField('y', 4, PointField.FLOAT32, 1), PointField('z', 8, PointField.FLOAT32, 1),] msg2.is_bigendian = False msg2.point_step = 12 #12 msg2.row_step = 12 * occ_coords.shape[0] msg2.is_dense = int(np.isfinite(occ_coords).all()) msg2.data = occ_coords.tostring() #np.asarray(occ_coords, np.float32).tostring() print("msg2: ", occ_coords.shape, len(msg2.data), msg2.row_step) self.occ_pub.publish(msg2) time4 = time.time() print("post pocess time: %fs"%(time4 - time3)) return
def __getitem__(self, idx): inputsdf_file = self.files[idx][0] sdf_file = self.files[idx][1] name = os.path.splitext(os.path.basename(inputsdf_file))[0] color_file = None if self.is_chunks else os.path.splitext( sdf_file)[0] + '.colors' sdf, world2grid, known, colors = data_util.load_sdf( sdf_file, load_sparse=self.load_tgt_sparse, load_known=self.load_known and self.is_chunks, load_colors=True, color_file=color_file) if sdf is None: return {'name': None} if self.load_known and not self.is_chunks: file_info = os.path.split(sdf_file) prefix = file_info[0] + '-complete' if 'color' in file_info[ 0] else file_info[0] + '_scanned' pad_known = (3, 6, 6) if 'color' in file_info[0] else (3, 3, 3) known_file = os.path.join( prefix, os.path.splitext(file_info[1])[0] + '.knw') known_file = known_file.replace('_trunc32-complete', '-complete') known = data_util.load_known(known_file, pad_known=pad_known, scale_to_dims=sdf.shape) input_color_file = None if self.is_chunks else os.path.splitext( inputsdf_file)[0] + '.colors' input, _, _, _, input_colors = data_util.load_sdf( inputsdf_file, load_sparse=True, load_known=False, load_colors=True, color_file=input_color_file) if input is None: return {'name': None} if self.color_truncation > 0: locs = input[0][np.abs(input[1]) > self.color_truncation] # to mask out input_colors[locs[:, 0], locs[:, 1], locs[:, 2], :] = 0 max_input_dim = np.max(input[0], 0) if max_input_dim[0] >= sdf.shape[0] or max_input_dim[1] >= sdf.shape[ 1] or max_input_dim[2] >= sdf.shape[2]: mask = np.logical_and( input[0][:, 0] < sdf.shape[0], np.logical_and(input[0][:, 1] < sdf.shape[1], input[0][:, 2] < sdf.shape[2])) input[0] = input[0][mask] input[1] = input[1][mask] image_depth = None image_color = None image_pose = None image_intrinsic = None if self.frame_info_path and self.frame_path: image_depth, image_color, image_pose, image_intrinsic, image_frameids = data_util.load_frames( [name], [world2grid], self.frame_info_path, self.frame_path, randomize_frames=self.randomize_frames, depth_image_dims=self.image_dims, color_image_dims=self.image_dims, color_normalization=None, load_depth=self.load_depth, load_color=self.load_color) if image_color is not None: if self.load_depth: image_depth = image_depth.squeeze(1) if image_color is not None: if len(image_color.shape) == 3: image_color = image_color.squeeze(1) image_pose = image_pose.squeeze(1) image_intrinsic = image_intrinsic.squeeze(1) else: image_color = image_color.squeeze(0) image_pose = image_pose.squeeze(0) image_intrinsic = image_intrinsic.squeeze(0) if self.subsamp2d_factor > 1: sz = image_color.shape[2:] image_color = torch.nn.functional.interpolate( image_color, scale_factor=1.0 / self.subsamp2d_factor, mode='bilinear', align_corners=False) image_color = torch.nn.functional.interpolate( image_color, size=sz, mode='bilinear', align_corners=False) input_dense = data_util.sparse_to_dense_np(input[0], input[1][:, np.newaxis], sdf.shape[2], sdf.shape[1], sdf.shape[0], -float('inf')) if (self.is_chunks and (self.input_dim[0] != 96 and self.input_dim[0] != 128 and self.input_dim[0] != 160)) or self.scene_subsample_factor > 1: scale_factor = float( self.input_dim[0]) / 128 if self.is_chunks else 1.0 / float( self.scene_subsample_factor) input_dense = torch.nn.functional.interpolate( torch.from_numpy(input_dense).unsqueeze(0).unsqueeze(0), scale_factor=scale_factor) * scale_factor input_dense = input_dense[0, 0].numpy() input_colors = torch.nn.functional.interpolate( torch.from_numpy(input_colors).permute( 3, 0, 1, 2).contiguous().unsqueeze(0).float(), scale_factor=scale_factor) input_colors = input_colors[0].permute( 1, 2, 3, 0).contiguous().numpy().astype(np.uint8) sdf = torch.nn.functional.interpolate( torch.from_numpy(sdf).unsqueeze(0).unsqueeze(0), scale_factor=scale_factor) * scale_factor sdf = sdf[0, 0].numpy() colors = torch.nn.functional.interpolate( torch.from_numpy(colors).permute( 3, 0, 1, 2).contiguous().unsqueeze(0).float(), scale_factor=scale_factor) colors = colors[0].permute(1, 2, 3, 0).contiguous().numpy().astype(np.uint8) if known is not None: known = torch.nn.functional.interpolate( torch.from_numpy(known).float().unsqueeze(0).unsqueeze(0), scale_factor=scale_factor).byte() known = known[0, 0].numpy() world2grid = np.matmul( data_util.make_scale_transform(scale_factor), world2grid).astype(np.float32) if self.augment_rgb_scaling: scale = np.random.rand(1) * ( self.aug_scale_range[1] - self.aug_scale_range[0]) + self.aug_scale_range[0] input_colors = data_util.convert_rgbgrid_to_hsvgrid( input_colors.astype(np.float32) / 255.0) colors = data_util.convert_rgbgrid_to_hsvgrid( colors.astype(np.float32) / 255.0) # hue augmentation scaled = input_colors[:, :, :, 0] * scale[0] cmask = scaled >= 360 if np.sum(cmask) > 0: scaled[cmask] = scaled[cmask] % 360 input_colors[:, :, :, 0] = scaled input_colors = np.clip( data_util.convert_hsvgrid_to_rgbgrid(input_colors) * 255, 0, 255).astype(np.uint8) scaled = colors[:, :, :, 0] * scale[0] cmask = scaled >= 360 if np.sum(cmask) > 0: scaled[cmask] = scaled[cmask] % 360 colors[:, :, :, 0] = scaled colors = np.clip( data_util.convert_hsvgrid_to_rgbgrid(colors) * 255, 0, 255).astype(np.uint8) if image_color is not None: image_color = data_util.convert_rgbgrid_to_hsvgrid( image_color.permute(0, 2, 3, 1).numpy()) scaled = image_color[:, :, :, 0] * scale[0] cmask = scaled >= 360 if np.sum(cmask) > 0: scaled[cmask] = scaled[cmask] % 360 image_color[:, :, :, 0] = scaled image_color = data_util.convert_hsvgrid_to_rgbgrid(image_color) image_color = torch.from_numpy(image_color).permute( 0, 3, 1, 2).contiguous() if self.color_space == 'lab': colors = skcolor.rgb2lab(colors.astype(np.float32) / 255).astype( np.float32) input_colors = skcolor.rgb2lab( input_colors.astype(np.float32) / 255).astype(np.float32) # normalize tgt to 255 (expected in loss) colors[:, :, :, 0] = (colors[:, :, :, 0] / 100.0) * 255.0 colors[:, :, :, 1:] = (colors[:, :, :, 1:] + 100.0) / 200.0 * 255.0 # normalize input to 0/255 input_colors[:, :, :, 0] = input_colors[:, :, :, 0] / 100.0 input_colors[:, :, :, 1:] = (input_colors[:, :, :, 1:] + 100.0) / 200.0 input_colors *= 255.0 # images if image_color is not None: image_color = image_color.permute(0, 2, 3, 1).contiguous().view( 1, -1, 3).numpy() image_color = skcolor.rgb2lab(image_color).astype(np.float32) image_color = np.transpose(image_color, [0, 2, 1]).reshape( 1, 3, self.image_dims[1], self.image_dims[0]) # normalize to 0/1 image_color[:, 0] = image_color[:, 0] / 100.0 image_color[:, 1:] = (image_color[:, 1:] + 100.0) / 200.0 image_color = torch.from_numpy(image_color.astype(np.float32)) empty = np.abs(input_dense) > self.truncation mask = np.zeros(input_dense.shape, dtype=np.float32) mask[input_dense <= -1] = 1 mask[empty] = 0 input_dense = data_util.preprocess_sdf_pt(input_dense, self.truncation) input_colors = input_colors.astype(np.float32) / 255.0 input_colors[empty] = 0 colors = torch.from_numpy(colors) input_dense = torch.from_numpy(input_dense).unsqueeze(0) input_colors = torch.from_numpy(input_colors).permute(3, 0, 1, 2).contiguous() input = torch.cat([input_dense, input_colors], 0) mask = torch.from_numpy(mask).unsqueeze(0) sdf = sdf[np.newaxis, :] sdf = torch.from_numpy(sdf) world2grid = torch.from_numpy(world2grid) if self.load_known: known = known[np.newaxis, :] known = torch.from_numpy(known) sample = { 'name': name, 'input': input, 'sdf': sdf, 'world2grid': world2grid, 'known': known, 'colors': colors, 'image_depth': image_depth, 'image_color': image_color, 'image_pose': image_pose, 'image_intrinsic': image_intrinsic, 'mask': mask } return sample
def tsdf_callback(data): # gen = point_cloud2.read_points(data, skip_nans=True) # print(type(gen)) # exit() time1 = time.time() gen = point_cloud2.read_points(data.points, field_names=("x", "y", "z", "intensity"), skip_nans=True) center = data.center points = np.array(list(gen)).reshape([-1, 4]) print(points.shape, points[:, -1].max()) offset_x = center[0] - args.voxel_size * args.dimx / 2 offset_y = center[1] - args.voxel_size * args.dimx / 2 offset_z = -2.7 # points[:, 2].min() -0.1 points[:, 0] -= offset_x points[:, 1] -= offset_y points[:, 2] -= offset_z # print("minz: ", offset_z) points = points / args.voxel_size locs = np.floor(points[:, :3]).astype(np.int32) feats = points[:, 3] # print("feats = ",feats) mask = (locs[:, 0] < args.dimx) & (locs[:, 1] < args.dimx) & ( locs[:, 2] < args.dimz) & (locs[:, 0] >= 0) & (locs[:, 1] >= 0) & (locs[:, 2] >= 0) locs = locs[mask] feats = feats[mask] input = data_util.sparse_to_dense_np(locs, feats, args.dimx, args.dimx, args.dimz, -float('inf')) input = torch.from_numpy(input) #.unsqueeze(0) # print("shape after s to d", input.shape) original_input = input print("original_input: ", original_input.shape) # for grids whose tsdf value > -2, we consider it as already known known_mask = (original_input > -2) original_occ = (original_input.abs() < args.occ_thresord) original_free = original_input > args.occ_thresord original_unkown = original_input < -2 # print("INPUT: %f occ, %f free, %f unkwn"%(original_occ.float().sum()/SUM, original_free.float().sum()/SUM, original_unkown.float().sum()/SUM)) if args.flipped: input[input.abs() > args.truncation] = 0 input = torch.sign(input) * (args.truncation - input.abs()) else: input[input > args.truncation] = args.truncation input[input < -args.truncation] = -args.truncation # print("input mean:",input.abs().mean(), input.mean(), " max=",input.max()) # print("known: %f %%"%(known_mask.float().sum() * 100 /(64*64*48))) time2 = time.time() print("prepocess time: %fs" % (time2 - time1)) with torch.no_grad(): # print("input_dim: ", np.array(locs.shape)) input = input.unsqueeze(0) if not args.cpu: input = input.cuda() output_occs = None # try: output_sdf, output_occs = model(input, WEIGHTS) # print("out mean: ", output_sdf.abs().mean()) # except: # print("MODEL FAIL") # # embed() # output_sdf = output_sdf/args.voxel_size # fix conflicts with input in known grids known_mask = known_mask.unsqueeze(0).unsqueeze(0) output_sdf[known_mask] = input[known_mask] if not args.flipped: sdf_to_occ = output_sdf.abs( ) < args.occ_thresord # < n (n>0, meaner when n get smaller) else: sdf_to_occ = output_sdf.abs( ) > args.occ_thresord # > n (n<3, meaner when n get bigger) sdf_to_occ = sdf_to_occ.cpu() # wlz: original occ:input, sdf_to_occ:output, added_occ:added sdf_to_occ = sdf_to_occ.squeeze(0).squeeze(0) # sdf_to_occ = output_occs[-1][0].cpu().squeeze(0) # sdf_to_occ = torch.sigmoid(sdf_to_occ) > 0.5 added_occ = sdf_to_occ.float() > original_occ.float() time3 = time.time() print("model time: %fs" % (time3 - time2)) if sdf_to_occ is not None: # # output_occs = output_occs.cpu().squeeze() # original_occ = original_occ.squeeze() # occ_coords = POINTER[original_occ].reshape([-1, 3]) # num_points = occ_coords.shape[0] # added_color = torch.ones(num_points).unsqueeze(1).cuda() * (-0.2) # occ_coords = torch.cat((occ_coords, added_color), 1) # print("original occ size: ",occ_coords.shape) added_occ = added_occ.squeeze() occ_coords_1 = POINTER[added_occ].reshape([-1, 3]) num_points = occ_coords_1.shape[0] added_color_1 = torch.ones(num_points).unsqueeze(1).cuda() * 10 occ_coords_1 = torch.cat((occ_coords_1, added_color_1), 1) # print("added size: ",occ_coords_1.shape) # occ_coords = torch.cat((occ_coords, occ_coords_1), 0) occ_coords = occ_coords_1 if occ_coords.shape[0] > 0: occ_coords *= args.voxel_size occ_coords[:, 0] += offset_x occ_coords[:, 1] += offset_y occ_coords[:, 2] += offset_z occ_coords = occ_coords.cpu().numpy() msg = data.points msg.height = 1 msg.width = occ_coords.shape[0] msg.fields = [ PointField('x', 0, PointField.FLOAT32, 1), PointField('y', 4, PointField.FLOAT32, 1), PointField('z', 8, PointField.FLOAT32, 1), PointField('intensity', 12, PointField.FLOAT32, 1) ] msg.is_bigendian = False msg.point_step = 16 #12 msg.row_step = 16 * occ_coords.shape[0] msg.is_dense = int(np.isfinite(occ_coords).all()) msg.data = occ_coords.tostring( ) #np.asarray(occ_coords, np.float32).tostring() # print(occ_coords.shape, len(msg.data), msg.row_step) # print("postpocess time: %fs"%(time.time() - time3)) publish_pcl(msg) #publish added occ # original_unkown = original_input <= -2 original_unkown = original_unkown.squeeze() occ_coords = POINTER[original_occ].reshape([-1, 3]) # print("known occ size: ",occ_coords.shape) if occ_coords.shape[0] > 0: occ_coords *= args.voxel_size occ_coords[:, 0] += offset_x occ_coords[:, 1] += offset_y occ_coords[:, 2] += offset_z occ_coords = occ_coords.cpu().numpy() msg = data.points msg.height = 1 msg.width = occ_coords.shape[0] msg.fields = [ PointField('x', 0, PointField.FLOAT32, 1), PointField('y', 4, PointField.FLOAT32, 1), PointField('z', 8, PointField.FLOAT32, 1), ] msg.is_bigendian = False msg.point_step = 12 #12 msg.row_step = 12 * occ_coords.shape[0] msg.is_dense = int(np.isfinite(occ_coords).all()) msg.data = occ_coords.tostring( ) #np.asarray(occ_coords, np.float32).tostring() publish_added_pcl(msg) return