depth_mapping_3d = np.ones(vox_size[2] * vox_size[1] * vox_size[0], dtype=np.float32) * (-1) dp.TSDF((depth / 1000.0).reshape(-1), cam_K, cam_pose, vox_origin, vox_unit, vox_size.astype(np.float32), vox_margin, depth.shape[0], depth.shape[1], tsdf, depth_mapping_3d) depth_mapping_3d = depth_mapping_3d.astype(np.int64) label = np.zeros(vox_size[2] * vox_size[1] * vox_size[0] / (sampleRatio * sampleRatio * sampleRatio), dtype=np.float32) tsdf_downsample = np.zeros(vox_size[2] * vox_size[1] * vox_size[0] / (sampleRatio * sampleRatio * sampleRatio), dtype=np.float32) dp.DownSampleLabel(readin_bin, vox_size.astype(np.float32), sampleRatio, tsdf, label, tsdf_downsample) label = label.astype(np.int32) label_weight = getLabelWeight(label, tsdf_downsample) label[np.where(label == 255)] = 0 dp.TSDFTransform(tsdf, vox_size.astype(np.float32)) tsdf = torch.FloatTensor(tsdf).view( (1, vox_size[2], vox_size[1], vox_size[0])) label = torch.LongTensor(label).view(vox_size[2] / sampleRatio, vox_size[1] / sampleRatio, vox_size[0] / sampleRatio) label_weight = torch.FloatTensor(label_weight).view( vox_size[2] / sampleRatio, vox_size[1] / sampleRatio, vox_size[0] / sampleRatio) depth_mapping_3d = torch.LongTensor(depth_mapping_3d) np.savez_compressed(os.path.join(outputpath, '%06d.npz' % fileindex), tsdf.numpy(), label.numpy(), label_weight.numpy(),
def multi_process(index, lenlist, startindex): print str(index) + '/' + str(lenlist) binfile = open(binlist[index - startindex], 'rb') vox_origin = np.array(struct.unpack('3f', binfile.read(12)), dtype=np.float32) cam_pose = np.array(struct.unpack('16f', binfile.read(64)), dtype=np.float32) readin_bin = np.zeros(vox_size[0] * vox_size[1] * vox_size[2], dtype=np.float32) all_the_byte = binfile.read() binfile.close() all_the_val = struct.unpack("%dI" % (len(all_the_byte) / 4), all_the_byte) last_index = 0 count_bad_data = 0 for i in xrange(len(all_the_val) / 2): if last_index + all_the_val[ i * 2 + 1] > vox_size[0] * vox_size[1] * vox_size[2]: count_bad_data = 0 break if all_the_val[i * 2] == 255: readin_bin[last_index:last_index + all_the_val[i * 2 + 1]] = 255 else: readin_bin[last_index:last_index + all_the_val[i * 2 + 1]] = segmentation_class_map[ all_the_val[i * 2] % 37] if all_the_val[i * 2] != 0: count_bad_data += all_the_val[i * 2 + 1] last_index += all_the_val[i * 2 + 1] if count_bad_data < 10 or last_index != vox_size[0] * vox_size[ 1] * vox_size[2]: print str(index) + ' is awful!' return depth = cv2.imread(depthlist[index - startindex], -1) depth = ((depth << 13) | (depth >> 3)).astype(np.float32) tsdf = np.ones(vox_size[2] * vox_size[1] * vox_size[0], dtype=np.float32) # modified: zeros to ones depth_mapping_3d = np.ones(640 * 480, dtype=np.float32) * ( -1 ) # vox_size[2]*vox_size[1]*vox_size[0] seems to be wrong, it should be 640x480 dp.TSDF((depth / 1000.0).reshape(-1), cam_K, cam_pose, vox_origin, vox_unit, vox_size.astype(np.float32), vox_margin, depth.shape[0], depth.shape[1], tsdf, depth_mapping_3d) depth_mapping_3d = depth_mapping_3d.astype(np.int64) label = np.zeros(vox_size[2] * vox_size[1] * vox_size[0] / (sampleRatio * sampleRatio * sampleRatio), dtype=np.float32) tsdf_downsample = np.zeros(vox_size[2] * vox_size[1] * vox_size[0] / (sampleRatio * sampleRatio * sampleRatio), dtype=np.float32) dp.DownSampleLabel(readin_bin, vox_size.astype(np.float32), sampleRatio, tsdf, label, tsdf_downsample) label = label.astype(np.int32) label_weight = getLabelWeight(label, tsdf_downsample) label[np.where(label == 255)] = 0 dp.TSDFTransform(tsdf, vox_size.astype(np.float32)) tsdf = torch.FloatTensor(tsdf).view( (1, vox_size[2], vox_size[1], vox_size[0])) label = torch.LongTensor(label).view(vox_size[2] / sampleRatio, vox_size[1] / sampleRatio, vox_size[0] / sampleRatio) label_weight = torch.FloatTensor(label_weight).view( vox_size[2] / sampleRatio, vox_size[1] / sampleRatio, vox_size[0] / sampleRatio) depth_mapping_3d = torch.LongTensor(depth_mapping_3d) np.savez_compressed(os.path.join(outputpath, '%06d.npz' % index), tsdf.numpy(), label.numpy(), label_weight.numpy(), depth_mapping_3d.numpy())
def __getitem__(self, index): while True: binfile = open(self.binlist[index], 'rb') vox_origin = np.array(struct.unpack('3f', binfile.read(12)), dtype=np.float32) cam_pose = np.array(struct.unpack('16f', binfile.read(64)), dtype=np.float32) readin_bin = np.zeros(self.vox_size[0] * self.vox_size[1] * self.vox_size[2], dtype=np.float32) all_the_byte = binfile.read() binfile.close() all_the_val = struct.unpack("%dI" % (len(all_the_byte) / 4), all_the_byte) last_index = 0 count_bad_data = 0 for i in xrange(len(all_the_val) / 2): if last_index + all_the_val[i * 2 + 1] > self.vox_size[ 0] * self.vox_size[1] * self.vox_size[2]: count_bad_data = 0 break if all_the_val[i * 2] == 255: readin_bin[last_index:last_index + all_the_val[i * 2 + 1]] = 255 else: readin_bin[last_index:last_index + all_the_val[i * 2 + 1]] = self.segmentation_class_map[ all_the_val[i * 2] % 37] if all_the_val[i * 2] != 0: count_bad_data += all_the_val[i * 2 + 1] last_index += all_the_val[i * 2 + 1] if count_bad_data < 10 or last_index != self.vox_size[ 0] * self.vox_size[1] * self.vox_size[2]: index = self.hardexamples[random.randint( 0, len(self.hardexamples) - 1)] continue color = cv2.imread(self.colorlist[index], -1) depth = cv2.imread(self.depthlist[index], -1).astype(np.float32) tsdf = np.zeros(self.vox_size[2] * self.vox_size[1] * self.vox_size[0], dtype=np.float32) depth_mapping_3d = np.ones( self.vox_size[2] * self.vox_size[1] * self.vox_size[0], dtype=np.float32) * (-1) dp.TSDF((depth / 1000.0).reshape(-1), self.cam_K, cam_pose, vox_origin, self.vox_unit, self.vox_size.astype(np.float32), self.vox_margin, depth.shape[0], depth.shape[1], tsdf, depth_mapping_3d) depth_mapping_3d = depth_mapping_3d.astype(np.int64) label = np.zeros( self.vox_size[2] * self.vox_size[1] * self.vox_size[0] / (self.sampleRatio * self.sampleRatio * self.sampleRatio), dtype=np.float32) tsdf_downsample = np.zeros( self.vox_size[2] * self.vox_size[1] * self.vox_size[0] / (self.sampleRatio * self.sampleRatio * self.sampleRatio), dtype=np.float32) dp.DownSampleLabel(readin_bin, self.vox_size.astype(np.float32), self.sampleRatio, tsdf, label, tsdf_downsample) label = label.astype(np.int32) label_weight = getLabelWeight(label, tsdf_downsample) label[np.where(label == 255)] = 0 dp.TSDFTransform(tsdf, self.vox_size.astype(np.float32)) color = np.transpose(color, (2, 0, 1)).astype(np.float32) / 255.0 depth1 = torch.FloatTensor(depth / 10000.0).view( 1, depth.shape[0], depth.shape[1]) color = torch.FloatTensor(color) color = self.img_transform(color) img = torch.cat((color, depth1), dim=0).contiguous() tsdf = torch.FloatTensor(tsdf).view( (1, self.vox_size[2], self.vox_size[1], self.vox_size[0])) label = torch.LongTensor(label).view( self.vox_size[2] / self.sampleRatio, self.vox_size[1] / self.sampleRatio, self.vox_size[0] / self.sampleRatio) label_weight = torch.FloatTensor(label_weight).view( self.vox_size[2] / self.sampleRatio, self.vox_size[1] / self.sampleRatio, self.vox_size[0] / self.sampleRatio) depth_mapping_3d = torch.LongTensor(depth_mapping_3d) return img, tsdf, label, label_weight, depth_mapping_3d