S23DR-P2R / format_dataset.py
colin1842's picture
add model
8d5039c
import os
import shutil
import logging
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def create_directory(path):
if not os.path.exists(path):
os.makedirs(path)
def transfer_files(source_clean_xyz_dir, source_gt_dir, destination_dir, txt_file_path):
create_directory(destination_dir)
subdirectory_paths = []
for filename in os.listdir(source_clean_xyz_dir):
if filename.endswith('.xyz'):
base_name = os.path.splitext(filename)[0]
new_subdir = os.path.join(destination_dir, base_name)
create_directory(new_subdir)
subdirectory_paths.append(new_subdir)
source_xyz = os.path.join(source_clean_xyz_dir, filename)
destination_xyz = os.path.join(new_subdir, 'points.xyz')
shutil.copy(source_xyz, destination_xyz)
logging.info(f'Copied {source_xyz} to {destination_xyz}')
source_obj = os.path.join(source_gt_dir, f'{base_name}.obj')
destination_obj = os.path.join(new_subdir, 'polygon.obj')
if os.path.exists(source_obj):
shutil.copy(source_obj, destination_obj)
logging.info(f'Copied {source_obj} to {destination_obj}')
else:
logging.warning(f'File not found: {source_obj}')
with open(txt_file_path, 'w') as txt_file:
for path in subdirectory_paths:
txt_file.write(path + '\n')
logging.info(f'Written subdirectory paths to {txt_file_path}')
# Define paths
source_clean_xyz_dir = 'Data/hoho_data_train/clean_xyz'
source_gt_dir = 'Data/hoho_data_train/gt'
destination_dir = 'Data/hoho_data_train'
txt_file_path = 'Data/hoho_data_train/train.txt'
# Run the transfer process
transfer_files(source_clean_xyz_dir, source_gt_dir, destination_dir, txt_file_path)