1960-utils/Texturing/MergeTextures2/merge_textures.py

import os
import sys
from PIL import Image # type: ignore
import time

# Define suffix lists for BaseColor, Normal, RMA/ORM
BASECOLOR_SUFFIXES = ['_alb.', '_albedo.', '_bc.', '_basecolor.', '_b.']
NORMAL_SUFFIXES = ['_nrm.', '_normal.', '_n.']
RMA_SUFFIXES = ['_rma.']
ORM_SUFFIXES = ['_orm.']
EMISSIVE_SUFFIXES = ['_emissive.']
OPACITY_SUFFIXES = ['_opacity.']
MASK_SUFFIXES = ['_mask.','_m.']

def detect_texture_type(filename):
    """ Detects the type of texture based on its suffix """
    if any(suffix in filename.lower() for suffix in BASECOLOR_SUFFIXES):
        return 'BaseColor'
    elif any(suffix in filename.lower() for suffix in NORMAL_SUFFIXES):
        return 'Normal'
    elif any(suffix in filename.lower() for suffix in RMA_SUFFIXES):
        return 'RMA'
    elif any(suffix in filename.lower() for suffix in ORM_SUFFIXES):
        return 'ORM'
    elif any(suffix in filename.lower() for suffix in EMISSIVE_SUFFIXES):
        return 'Emissive'
    elif any(suffix in filename.lower() for suffix in OPACITY_SUFFIXES):
        return 'Opacity'
    elif any(suffix in filename.lower() for suffix in MASK_SUFFIXES):
        return 'Mask'
    return None

def get_material_name(filename):
    """ Strips the 'T_' or 'TX_' prefix but keeps the suffix for texture type detection.
    Returns the full material name without the suffix for output file naming. """
    base_name = os.path.basename(filename)
    
    # Remove the 'T_' or 'TX_' prefix
    if base_name.startswith('T_'):
        base_name = base_name[2:]
    elif base_name.startswith('TX_'):
        base_name = base_name[3:]
    
    # Return the base_name without the suffix for output naming
    return base_name.rsplit('_', 1)[0]  # Split only at the last underscore

def process_textures(input_files):
    """ Main function to process all textures in a folder and convert to BCR/NMO """
    textures = {}
    
    # Group files by material name
    for filepath in input_files:
        filename = os.path.basename(filepath)
        material_name = get_material_name(filename)
        texture_type = detect_texture_type(filename)
        
        if material_name not in textures:
            textures[material_name] = {}
        textures[material_name][texture_type] = filepath
    
    # Create a merged folder in the same directory as the input
    base_path = os.path.dirname(input_files[0])
    output_folder = os.path.join(base_path, 'merged')
    os.makedirs(output_folder, exist_ok=True)
	
    material_count = len(textures)
    print(f"Detected {material_count} Materials to process.")

    failed_converts = 0
	
    # Process each material group
    for index, (material, files) in enumerate(textures.items()):
        basecolor_file = files.get('BaseColor')
        normal_file = files.get('Normal')
        rma_file = files.get('RMA')
        orm_file = files.get('ORM')
        emissive_file = files.get('Emissive')
        opacity_file = files.get('Opacity')
        mask_file = files.get('Mask')
		
        missing_files = []
		
		# Check for required textures
        if not basecolor_file:
            missing_files.append('BaseColor')
        if not normal_file:
            missing_files.append('Normal')
        if not (rma_file or orm_file):
            missing_files.append('RMA or ORM')
        
        # Report missing files if any
        if missing_files:
            print(f"({index + 1}/{material_count}) Skipping {material}: missing {', '.join(missing_files)}")
            failed_converts += 1
        else:
            # Convert to BCR/NMO format and track success or failure
            
            if convert_to_bcr_nmo(material, basecolor_file, normal_file, rma_file, orm_file, emissive_file, opacity_file, mask_file, output_folder):
                print(f"({index + 1}/{material_count}) {material}: Successfully converted.")
            else:
                failed_converts += 1  # Increment counter here if conversion fails
                print(f"({index + 1}/{material_count}) Skipping {material}: input file sizes do not match.")
	
    print(f"+++{material_count - failed_converts} of {material_count} materials successfully converted+++")
    time.sleep(3)

def convert_to_bcr_nmo(material, basecolor_file, normal_file, rma_file, orm_file, emissive_file, opacity_file, mask_file, output_folder):
    """ Converts given textures to BCR and NMO formats """
    basecolor_img = Image.open(basecolor_file).convert('RGBA')
    normal_img = Image.open(normal_file).convert('RGBA')
    
    if rma_file:
        rma_img = Image.open(rma_file).convert('RGBA')
        if not (basecolor_img.size == normal_img.size == rma_img.size):
            return False
        # BCR conversion
        bcr_img = Image.merge('RGBA', (basecolor_img.split()[0], basecolor_img.split()[1], basecolor_img.split()[2], rma_img.split()[0]))  # Use Roughness (Alpha from RMA/ORM)
        bcr_img.save(os.path.join(output_folder, f"{material}_BCR.png"))
        # NMO conversion
        nmo_img = Image.merge('RGBA', (normal_img.split()[0], normal_img.split()[1], rma_img.split()[1], rma_img.split()[2]))  # Use Metallic, AO from RMA/ORM
        nmo_img.save(os.path.join(output_folder, f"{material}_NMO.png"))
    elif orm_file:
        rma_img = Image.open(orm_file).convert('RGBA')
        if not (basecolor_img.size == normal_img.size == rma_img.size):
            return False
        # BCR conversion
        bcr_img = Image.merge('RGBA', (basecolor_img.split()[0], basecolor_img.split()[1], basecolor_img.split()[2], rma_img.split()[1]))  # Use Roughness (Alpha from RMA/ORM)
        bcr_img.save(os.path.join(output_folder, f"{material}_BCR.png"))
        # NMO conversion
        nmo_img = Image.merge('RGBA', (normal_img.split()[0], normal_img.split()[1], rma_img.split()[2], rma_img.split()[0]))  # Use Metallic, AO from RMA/ORM
        nmo_img.save(os.path.join(output_folder, f"{material}_NMO.png"))
    # Optionally handle emissive and opacity maps
    if emissive_file:
        emissive_img = Image.open(emissive_file).convert('RGB')
        emissive_img.save(os.path.join(output_folder, f"{material}_EM.png"))
    
    if opacity_file:
        opacity_img = Image.open(opacity_file).convert('L')
        opacity_img.save(os.path.join(output_folder, f"{material}_OP.png"))
    
    if mask_file:
        mask_img = Image.open(mask_file).convert('L')
        mask_img.save(os.path.join(output_folder, f"{material}_MASK.png"))

    return True
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: drag and drop texture files onto the script")
    else:
        # Get the file paths from sys.argv (ignoring the first argument which is the script name)
        input_files = sys.argv[1:]
        process_textures(input_files)