Function main

recompiler.py:46–93

()

    return parser.parse_args()

def main():
    args = parse_args()
    model_file = args.model

    if not model_file:
        print("No model file specified.")
        return

    if not os.path.isfile(model_file):
        print(f"The file {model_file} does not exist.")
        sys.exit(1)

    # Read the source code from the model file
    with open(model_file, 'r') as f:
        model_code = f.read()

    # Custom layer name
    custom_layer = 'MyDense'

    # Parse model source code to AST
    tree = ast.parse(model_code)

    # Transformer to replace Dense with MyDense
    class LayerReplacer(ast.NodeTransformer):
        def visit_Call(self, node):
            if isinstance(node.func, ast.Name) and node.func.id == 'Dense':
                node.func.id = custom_layer  
            return self.generic_visit(node)

    transformer = LayerReplacer()
    modified_tree = transformer.visit(tree)

    # Generate modified source code
    modified_code = ast.unparse(modified_tree)

    print("Modified Model:\n", modified_code)

    output_file = f'encrypt_{model_file}'

    # Write the modified model code to a new file
    try:
        with open(output_file, 'w') as file:
            file.write(modified_code)
            file.write("\n\n" + MyDense_complate)
        
        print(f"File {output_file} has been created with the modified code.")
    except IOError as e:
        print(f"Failed to write to file {output_file}: {e}")

if __name__ == "__main__":
    main()