因为官方给的reader.py只能一个个生成太麻烦了，于是写了个小py节省时间。。
官方代码来源：https://github.com/ScanNet/ScanNet/tree/master/SensReader/python
官方的reader.py文件
import argparse import os, sys from SensorData import SensorData # params parser = argparse.ArgumentParser() # data paths parser.add_argument('--filename', required=True, help='path to sens file to read') parser.add_argument('--output_path', required=True, help='path to output folder') parser.add_argument('--export_depth_images', dest='export_depth_images', action='store_true') parser.add_argument('--export_color_images', dest='export_color_images', action='store_true') parser.add_argument('--export_poses', dest='export_poses', action='store_true') parser.add_argument('--export_intrinsics', dest='export_intrinsics', action='store_true') parser.set_defaults(export_depth_images=True, export_color_images=True, export_poses=True, export_intrinsics=True) opt = parser.parse_args() print(opt) def main(): if not os.path.exists(opt.output_path): os.makedirs(opt.output_path) # load the data sys.stdout.write('loading %s...' % opt.filename) sd = SensorData(opt.filename) sys.stdout.write('loaded!\n') if opt.export_depth_images: sd.export_depth_images(os.path.join(opt.output_path, 'depth')) if opt.export_color_images: sd.export_color_images(os.path.join(opt.output_path, 'color')) if opt.export_poses: sd.export_poses(os.path.join(opt.output_path, 'pose')) if opt.export_intrinsics: sd.export_intrinsics(os.path.join(opt.output_path, 'intrinsic')) if __name__ == '__main__': main()
官方的SensorData文件
from tqdm import tqdm import os, struct import numpy as np import zlib import imageio import cv2 import png COMPRESSION_TYPE_COLOR = {-1:'unknown', 0:'raw', 1:'png', 2:'jpeg'} COMPRESSION_TYPE_DEPTH = {-1:'unknown', 0:'raw_ushort', 1:'zlib_ushort', 2:'occi_ushort'} class RGBDFrame(): def load(self, file_handle): self.camera_to_world = np.asarray(struct.unpack('f'*16, file_handle.read(16*4)), dtype=np.float32).reshape(4, 4) self.timestamp_color = struct.unpack('Q', file_handle.read(8))[0] self.timestamp_depth = struct.unpack('Q', file_handle.read(8))[0] self.color_size_bytes = struct.unpack('Q', file_handle.read(8))[0] self.depth_size_bytes = struct.unpack('Q', file_handle.read(8))[0] # self.color_data = ''.join(struct.unpack('c'*self.color_size_bytes, file_handle.read(self.color_size_bytes))) # self.depth_data = ''.join(struct.unpack('c'*self.depth_size_bytes, file_handle.read(self.depth_size_bytes))) self.color_data = b''.join(struct.unpack('c' * self.color_size_bytes, file_handle.read(self.color_size_bytes))) self.depth_data = b''.join(struct.unpack('c' * self.depth_size_bytes, file_handle.read(self.depth_size_bytes))) def decompress_depth(self, compression_type): if compression_type == 'zlib_ushort': return self.decompress_depth_zlib() else: raise def decompress_depth_zlib(self): return zlib.decompress(self.depth_data) def decompress_color(self, compression_type): if compression_type == 'jpeg': return self.decompress_color_jpeg() else: raise def decompress_color_jpeg(self): return imageio.imread(self.color_data) class SensorData: def __init__(self, filename): self.version = 4 self.load(filename) def load(self, filename): with open(filename, 'rb') as f: version = struct.unpack('I', f.read(4))[0] assert self.version == version strlen = struct.unpack('Q', f.read(8))[0] # self.sensor_name = ''.join(struct.unpack('c'*strlen, f.read(strlen))) self.sensor_name = b''.join(struct.unpack('c'*strlen, f.read(strlen))) self.intrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) self.extrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) self.intrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) self.extrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4) self.color_compression_type = COMPRESSION_TYPE_COLOR[struct.unpack('i', f.read(4))[0]] self.depth_compression_type = COMPRESSION_TYPE_DEPTH[struct.unpack('i', f.read(4))[0]] self.color_width = struct.unpack('I', f.read(4))[0] self.color_height = struct.unpack('I', f.read(4))[0] self.depth_width = struct.unpack('I', f.read(4))[0] self.depth_height = struct.unpack('I', f.read(4))[0] self.depth_shift = struct.unpack('f', f.read(4))[0] num_frames = struct.unpack('Q', f.read(8))[0] self.frames = [] for i in tqdm(range(num_frames), ncols=80): frame = RGBDFrame() frame.load(f) self.frames.append(frame) def export_depth_images(self, output_path, image_size=None, frame_skip=1): if not os.path.exists(output_path): os.makedirs(output_path) print ('exporting', len(self.frames)//frame_skip, ' depth frames to', output_path) for f in tqdm(range(0, len(self.frames), frame_skip), ncols=80): depth_data = self.frames[f].decompress_depth(self.depth_compression_type) depth = np.fromstring(depth_data, dtype=np.uint16).reshape(self.depth_height, self.depth_width) if image_size is not None: depth = cv2.resize(depth, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST) #imageio.imwrite(os.path.join(output_path, str(f) + '.png'), depth) with open(os.path.join(output_path, str(f) + '.png'), 'wb') as f: # write 16-bit writer = png.Writer(width=depth.shape[1], height=depth.shape[0], bitdepth=16) depth = depth.reshape(-1, depth.shape[1]).tolist() writer.write(f, depth) def export_color_images(self, output_path, image_size=None, frame_skip=1): if not os.path.exists(output_path): os.makedirs(output_path) print('exporting', len(self.frames) // frame_skip, 'color frames to', output_path) for f in tqdm(range(0, len(self.frames), frame_skip), ncols=80): color = self.frames[f].decompress_color(self.color_compression_type) if image_size is not None: color = cv2.resize(color, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST) imageio.imwrite(os.path.join(output_path, str(f) + '.jpg'), color) def save_mat_to_file(self, matrix, filename): with open(filename, 'w') as f: for line in matrix: np.savetxt(f, line[np.newaxis], fmt='%f') def export_poses(self, output_path, frame_skip=1): if not os.path.exists(output_path): os.makedirs(output_path) print('exporting', len(self.frames) // frame_skip, 'camera poses to', output_path) for f in range(0, len(self.frames), frame_skip): self.save_mat_to_file(self.frames[f].camera_to_world, os.path.join(output_path, str(f) + '.txt')) def export_intrinsics(self, output_path): if not os.path.exists(output_path): os.makedirs(output_path) print('exporting camera intrinsics to', output_path) self.save_mat_to_file(self.intrinsic_color, os.path.join(output_path, 'intrinsic_color.txt')) self.save_mat_to_file(self.extrinsic_color, os.path.join(output_path, 'extrinsic_color.txt')) self.save_mat_to_file(self.intrinsic_depth, os.path.join(output_path, 'intrinsic_depth.txt')) self.save_mat_to_file(self.extrinsic_depth, os.path.join(output_path, 'extrinsic_depth.txt'))
然后本人写的自动导出数据文件，直接替换reader.py就行了
记得修改文件夹路径。。。。
import argparse import os, sys import glob from SensorData import SensorData filename = '' output_path = '' # params parser = argparse.ArgumentParser() # data paths parser.add_argument('--filename', required=False, help='path to sens file to read', default=filename) parser.add_argument('--output_path', required=False, help='path to output folder', default=output_path) parser.add_argument('--export_depth_images', dest='export_depth_images', action='store_true') parser.add_argument('--export_color_images', dest='export_color_images', action='store_true') parser.add_argument('--export_poses', dest='export_poses', action='store_true') parser.add_argument('--export_intrinsics', dest='export_intrinsics', action='store_true') parser.set_defaults(export_depth_images=True, export_color_images=True, export_poses=True, export_intrinsics=True) opt = parser.parse_args() print(opt) def main(): if not os.path.exists(opt.output_path): os.makedirs(opt.output_path) # load the data sys.stdout.write('loading %s...' % opt.filename) sd = SensorData(opt.filename) sys.stdout.write('loaded!\n') if opt.export_depth_images: sd.export_depth_images(os.path.join(opt.output_path, 'depth')) if opt.export_color_images: sd.export_color_images(os.path.join(opt.output_path, 'color')) if opt.export_poses: sd.export_poses(os.path.join(opt.output_path, 'pose')) if opt.export_intrinsics: sd.export_intrinsics(os.path.join(opt.output_path, 'intrinsic')) if __name__ == '__main__': # 获取所有文件夹的路径 folder_paths = glob.glob('/home/xxxx/dataset/ScanNet/scans/*') # 计算文件夹的数量 file_count = len([name for name in folder_paths if os.path.isdir(name)]) # 读取txt文档中的文件路径 with open('scannet_list.txt', 'r') as f: file_paths = [line.strip() for line in f if line.strip()] # # 确保txt文档中的文件数目与实际文件数目相匹配 # assert len(file_paths) == file_count, "文件数目不匹配" print(file_count) # 遍历文件并执行main函数 for i in range(file_count): # 获取文件的完整路径 relative_path = file_paths[i] sens_path = relative_path + '.sens' full_path = os.path.join('/home/xxxx/dataset/ScanNet/scans/', relative_path) # 更新opt.filename和opt.output_path opt.filename = os.path.join(full_path,sens_path) opt.output_path = full_path # 获取文件所在目录作为输出路径 # 检查opt.output_path路径下是否存在color文件夹 if os.path.isdir(os.path.join(opt.output_path, 'color')): print(f"跳过存在color文件夹的路径：{opt.output_path}") continue # 跳过此次循环 # 执行main函数 main()
代码中用到的scannet_list.txt在文章顶部最终效果图
扫一扫
4482
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
