File size: 2,352 Bytes
2ecc7ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import argparse
import os
import sys
import shutil
import subprocess
def run_inference(input_path, output_path):
# 1. 获取绝对路径
cwd = os.path.dirname(os.path.abspath(__file__)) # 当前脚本所在目录
python_exe = sys.executable # 当前环境的 python
inference_script = os.path.join(cwd, "inference_codeformer.py")
print(f"🚀 [CodeFormer] 启动中...")
# 2. 准备临时输入目录 (CodeFormer 喜欢读文件夹)
# 它的逻辑是：输入文件夹 -> 结果文件夹
temp_in = os.path.join(cwd, "temp_worker_input")
if os.path.exists(temp_in): shutil.rmtree(temp_in)
os.makedirs(temp_in, exist_ok=True)
# 复制图片进去
img_name = os.path.basename(input_path)
shutil.copy(input_path, os.path.join(temp_in, img_name))
# 3. 构造命令
# -w 0.5: 平衡画质和保真度
# --face_upsample: 这一步会把人脸贴回原图，这是我们需要的
cmd = [
python_exe,
inference_script,
"-w", "0.5",
"--input_path", temp_in,
"--face_upsample"
]
# 4. 执行命令
try:
# cwd=cwd 保证它能找到 weights 文件夹
subprocess.run(cmd, cwd=cwd, check=True)
# 5. 提取结果
# 默认结果路径: results/temp_worker_input_0.5/final_results/图片名.png
res_dir = os.path.join(cwd, "results", "temp_worker_input_0.5", "final_results")
res_file = os.path.join(res_dir, img_name)
if os.path.exists(res_file):
os.makedirs(os.path.dirname(output_path), exist_ok=True)
shutil.copy(res_file, output_path)
print(f"✅ CodeFormer 处理完成: {output_path}")
else:
print(f"❌ 未找到结果文件: {res_file}")
print(" (可能是因为图片里没有人脸，或者模型下载失败)")
except Exception as e:
print(f"❌ 执行出错: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', required=True)
parser.add_argument('-o', '--output', required=True)
parser.add_argument('-m', '--model', help="Ignored")
args = parser.parse_args()
run_inference(args.input, args.output) |