Hugging Face's logo

hugaagg
/
IR_expeiment

Diffusers
Model card Files
xet
Community
IR_expeiment
File size: 5,654 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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
 
import os
import shutil
import subprocess
import glob

class AgentToolbox:
    def __init__(self):
        self.root_dir = r"G:\IR_Experiment"
        self.output_base = os.path.join(self.root_dir, "Agent_Workspace")
        os.makedirs(self.output_base, exist_ok=True)

        # ================= 环境配置 (关键！) =================
        # 请根据你的实际情况修改这里的 python.exe 路径
        # 1. ir_final 环境 (用于 DarkIR, NAFNet, PromptIR)
        self.env_main = r"D:\conda\envs\ir_final\python.exe"
        # 2. swinir_env 环境 (用于 SwinIR)
        self.env_swinir = r"D:\conda\envs\swinir_env\python.exe" 
        
        # 如果找不到路径，尝试用系统默认的 'python' (前提是你激活了对应环境)
        if not os.path.exists(self.env_main): self.env_main = "python"
        if not os.path.exists(self.env_swinir): self.env_swinir = "python"

    def _run_cmd(self, cmd, cwd):
        """执行命令行的通用函数"""
        print(f"\n[Toolbox] 正在执行: {cmd} ...")
        try:
            subprocess.run(cmd, shell=True, check=True, cwd=cwd)
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ 执行失败: {e}")
            return False

    # ================= 工具 1: DarkIR (低光增强) =================
    def call_darkir(self, image_path):
        print(f"🌙 [调用 DarkIR] 正在处理低光: {os.path.basename(image_path)}")
        tool_dir = os.path.join(self.root_dir, "DarkIR")
        
        # 1. 适配输入：DarkIR 也是读文件夹的，我们把图复制到它的 input 目录
        input_dir = os.path.join(tool_dir, "test_input") # 对应 run_darkir.py 里的路径
        if os.path.exists(input_dir): shutil.rmtree(input_dir)
        os.makedirs(input_dir, exist_ok=True)
        shutil.copy(image_path, os.path.join(input_dir, os.path.basename(image_path)))

        # 2. 调用我们之前写好的脚本
        # 注意：这里调用的是 run_darkir.py，确保它里面的路径是对的
        cmd = f'"{self.env_main}" run_darkir.py'
        
        if self._run_cmd(cmd, cwd=tool_dir):
            # 3. 提取输出
            # DarkIR 输出在 results 文件夹
            result_dir = os.path.join(tool_dir, "results")
            # 找到生成的文件
            res_files = glob.glob(os.path.join(result_dir, "*.*"))
            if res_files:
                # 把结果移动到 Agent 工作区
                out_name = f"darkir_{os.path.basename(image_path)}"
                final_path = os.path.join(self.output_base, out_name)
                shutil.copy(res_files[0], final_path)
                return final_path
        return None

    # ================= 工具 2: SwinIR (超分放大) =================
    def call_swinir(self, image_path, scale=4):
        print(f"🔍 [调用 SwinIR] 正在放大 {scale}倍: {os.path.basename(image_path)}")
        tool_dir = os.path.join(self.root_dir, "SwinIR")
        
        # 1. 适配输入
        temp_input = os.path.join(tool_dir, "testsets", "agent_temp")
        if os.path.exists(temp_input): shutil.rmtree(temp_input)
        os.makedirs(temp_input, exist_ok=True)
        shutil.copy(image_path, os.path.join(temp_input, os.path.basename(image_path)))

        # 2. 构建命令
        # 使用 Real-World x4 模型
        model_path = r"model_zoo/swinir/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth"
        cmd = f'"{self.env_swinir}" main_test_swinir.py --task real_sr --scale {scale} --model_path {model_path} --folder_lq testsets/agent_temp --tile 400'

        if self._run_cmd(cmd, cwd=tool_dir):
            # 3. 提取输出
            # SwinIR 结果通常在 results/swinir_real_sr_x4 里面
            result_dir = os.path.join(tool_dir, "results", f"swinir_real_sr_x{scale}")
            # 找最新生成的文件
            res_files = glob.glob(os.path.join(result_dir, "*.*"))
            if res_files:
                # SwinIR 会给文件名加后缀，我们找包含原名的那个
                target_file = [f for f in res_files if os.path.basename(image_path).split('.')[0] in f][-1]
                
                out_name = f"swinir_{os.path.basename(image_path)}"
                final_path = os.path.join(self.output_base, out_name)
                shutil.copy(target_file, final_path)
                return final_path
        return None

# ================= 模拟 Agent 调度逻辑 =================
if __name__ == "__main__":
    toolbox = AgentToolbox()
    
    # 1. 准备一张测试图 (你可以换成任何存在的图片路径)
    # 假设我们用之前生成的“低光”测试图
    original_img = r"G:\datasets\realblur_dataset_test\075_blur_1.png" # 确保这张图存在！
    
    if not os.path.exists(original_img):
        print("❌ 测试图不存在，请修改 original_img 路径")
        exit()

    print(f"🏁 开始处理任务: {original_img}")

    # --- 步骤 1: 先提亮 (DarkIR) ---
    bright_img = toolbox.call_darkir(original_img)
    
    if bright_img:
        print(f"✅ 第一步完成: {bright_img}")
        
        # --- 步骤 2: 再放大 (SwinIR) ---
        # 把第一步的结果喂给第二步
        final_img = toolbox.call_swinir(bright_img)
        
        if final_img:
            print(f"🎉 任务全部完成！最终结果: {final_img}")
        else:
            print("❌ 第二步 SwinIR 失败")
    else:
        print("❌ 第一步 DarkIR 失败")