import os

def patch_json_dump(target_file):
    if not os.path.exists(target_file):
        return
    with open(target_file, "r") as f:
        code = f.read()
        
    if "class NpEncoder(json.JSONEncoder):" in code:
        print(f"Already patched: {target_file}")
        return
        
    # 注入万能类型转换器
    encoder_code = """import json

class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if hasattr(obj, 'item'):
            return obj.item()
        return super().default(obj)
"""
    code = code.replace("import json", encoder_code)
    
    # 替换所有的 dump 调用
    code = code.replace("json.dump(stats, f, indent=4)", "json.dump(stats, f, indent=4, cls=NpEncoder)")
    code = code.replace("json.dump(metrics, f, indent=4)", "json.dump(metrics, f, indent=4, cls=NpEncoder)")
    code = code.replace("json.dump(results, f, indent=4)", "json.dump(results, f, indent=4, cls=NpEncoder)")
    
    with open(target_file, "w") as f:
        f.write(code)
    print(f"Patched JSON serialization in {target_file}")

patch_json_dump("/root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py")
patch_json_dump("/root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py")