import os
import subprocess

def convert_onnx_to_fp16_corrected(input_dir):
    # List all files in the input directory
    files = os.listdir(input_dir)
    
    # Filter out files with the .onnx extension
    onnx_files = [file for file in files if file.endswith('.onnx')]
    
    # Iterate over each ONNX file to convert it to FP16
    for onnx_file in onnx_files:
        # Split the file name to insert '.fp16' before the last underscore
        parts = onnx_file.rsplit('_', 1)
        if len(parts) == 2:
            # If there is at least one underscore, insert '.fp16' before the last part
            output_model_name = f"{parts[0]}.fp16_{parts[1]}"
        else:
            # If there's no underscore, just replace '.onnx' with '.fp16.onnx'
            output_model_name = onnx_file.replace('.onnx', '.fp16.onnx')
        
        # Construct the command to run the conversion script
        command = [
            'python3', 'convert_to_fp16.py',
            '--input_model', os.path.join(input_dir, onnx_file),
            '--output_model', os.path.join(input_dir, output_model_name)
        ]
        
        # Execute the command
        subprocess.run(command, check=True)

if __name__ == '__main__':
    # Assuming the current directory is the input directory
    input_dir = os.getcwd()
    convert_onnx_to_fp16_corrected(input_dir)