TRELLIS-Single3D / lut_extractor.py
gokaygokay's picture
pixelize
61554bf
raw
history blame
1.9 kB
import os
import zipfile
import shutil
from pathlib import Path
def clean_filename(filename):
# Split by '-' and keep only the part after it
if '-' in filename:
cleaned = filename.split('-')[1].strip()
# Ensure it ends with .cube
if cleaned.endswith('.cube'):
return cleaned
return filename
def extract_cube_files():
# Create directories if they don't exist
luts_dir = Path('luts')
cube_luts_dir = Path('cube_luts')
cube_luts_dir.mkdir(exist_ok=True)
# Process each zip file in the luts directory
for zip_path in luts_dir.glob('*.zip'):
try:
# Create a temporary directory for extraction
temp_dir = Path('temp_extract')
temp_dir.mkdir(exist_ok=True)
# Extract the zip file
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(temp_dir)
# Find and move all .cube files
for cube_file in temp_dir.rglob('*.cube'):
# Clean up the filename
new_name = clean_filename(cube_file.name)
dest_path = cube_luts_dir / new_name
# Handle duplicates
counter = 1
while dest_path.exists():
stem = Path(new_name).stem
suffix = Path(new_name).suffix
dest_path = cube_luts_dir / f"{stem}_{counter}{suffix}"
counter += 1
# Move the file
shutil.move(str(cube_file), str(dest_path))
# Clean up temporary directory
shutil.rmtree(temp_dir)
except Exception as e:
print(f"Error processing {zip_path}: {str(e)}")
continue
if __name__ == "__main__":
extract_cube_files()