Update app.py

app.py

@@ -15,13 +15,9 @@ from torch_geometric.data import Data as PyGData
 import matplotlib
 matplotlib.use('Agg')  # 修复后台线程问题
 from rdkit import Chem
-from rdkit.Chem import Draw
-from rdkit.Chem import AllChem
-from rdkit.Chem import MolFromSmiles
+from rdkit.Chem import Draw, AllChem, MolFromSmiles
 from io import BytesIO
 import traceback
-import base64
-import json
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -62,15 +58,10 @@ def load_models():
         "Brittle": ("models/best_model-B-6000-185.pth", 2)
     }
     models = {}
-    
     for name, (pth_path, output_dim) in model_info.items():
         logger.info(f"正在加载 {name} 模型: {pth_path}")
-        
-        # 检查模型文件是否存在
         if not os.path.exists(pth_path):
             logger.error(f"模型文件不存在: {pth_path}")
-            
-            # 尝试可能的文件名变体
             possible_files = [
                 pth_path,
                 pth_path.lower(),
@@ -78,7 +69,6 @@ def load_models():
                 pth_path.replace("-", "_"),
                 pth_path.replace("_", "-")
             ]
-            
             found = False
             for file in possible_files:
                 if os.path.exists(file):
@@ -86,119 +76,70 @@ def load_models():
                     pth_path = file
                     found = True
                     break
-            
             if not found:
-                logger.error(f"找不到任何匹配的模型文件: {pth_path}")
                 raise FileNotFoundError(f"模型文件 {pth_path} 不存在")
-            
         try:
-            # 修复模型初始化参数
             model = EnhancedGAT(input_dim=12, hidden_dim=512, output_dim=output_dim, num_heads=8)
-            
-            # 加载模型状态 - 解决 PyTorch 2.6+ 的安全问题
             logger.info(f"加载模型权重: {pth_path}")
-            
-            # 方法 1: 禁用 weights_only
             try:
-                # 尝试使用 weights_only=False 加载
                 state = torch.load(pth_path, map_location=device, weights_only=False)
-                logger.info("使用 weights_only=False 成功加载模型")
             except Exception as e:
-                logger.warning(f"使用 weights_only=False 加载失败: {str(e)}")
-                logger.info("尝试使用 weights_only=True 并添加安全全局变量")
-                
-                # 方法 2: 添加安全全局变量
+                logger.warning(f"weights_only=False 加载失败: {str(e)}")
                 try:
-                    # 导入必要的模块
-                    import numpy as np
-                    import torch.serialization
-                    
-                    # 添加安全全局变量
+                    import numpy as np, torch.serialization
                     torch.serialization.add_safe_globals([getattr(np, '_core', np).multiarray.scalar])
                     state = torch.load(pth_path, map_location=device, weights_only=True)
-                    logger.info("使用 weights_only=True 和安全全局变量成功加载模型")
                 except:
-                    # 最后尝试原始方式
-                    logger.warning("安全方式加载失败，尝试原始加载方式")
                     state = torch.load(pth_path, map_location=device)
-            
-            # 检查状态字典键名
             if "model_state_dict" in state:
                 state_dict = state["model_state_dict"]
-                logger.info("使用 'model_state_dict' 加载模型")
             else:
-                state_dict = state  # 直接使用整个文件
-                logger.info("使用整个状态字典加载模型")
-            
-            # 加载模型参数
+                state_dict = state
             model.load_state_dict(state_dict)
             model.eval().to(device)
             models[name] = model
             logger.info(f"{name} 模型加载成功!")
-            
         except Exception as e:
             logger.error(f"加载 {name} 模型失败: {str(e)}")
             raise
-    
     return models
 
 logger.info("开始加载所有模型...")
-try:
-    models = load_models()
-    logger.info("所有模型加载完成!")
-except Exception as e:
-    logger.error(f"模型加载过程中发生错误: {str(e)}")
-    raise
+models = load_models()
+logger.info("所有模型加载完成!")
 
 def predict_all(smiles):
     logger.info(f"收到预测请求: SMILES = {smiles}")
     try:
-        # 转换SMILES为图数据
-        logger.info("转换SMILES为图数据...")
         atom_features, (rows, cols, edge_attr), mol = smiles_to_graph(smiles)
-        
-        # 创建PyG数据对象
-        logger.info("创建PyG数据对象...")
         x = torch.tensor(atom_features, dtype=torch.float)
         edge_index = torch.tensor(np.column_stack((rows, cols)).T, dtype=torch.long)
         edge_attr = torch.tensor(edge_attr, dtype=torch.float).unsqueeze(1)
-        data = PyGData(
-            x=x,
-            edge_index=edge_index,
-            edge_attr=edge_attr,
-            smiles=[smiles],
-            batch=torch.zeros(x.size(0), dtype=torch.long)
-        )
-
+        data = PyGData(x=x,
+                       edge_index=edge_index,
+                       edge_attr=edge_attr,
+                       smiles=[smiles],
+                       batch=torch.zeros(x.size(0), dtype=torch.long))
         results = []
-        # 对每个模型进行预测
         for name in ["Elastic", "Plastic", "Brittle"]:
             logger.info(f"使用 {name} 模型进行预测...")
             try:
                 buf, pred = visualize_single_molecule(models[name], data, device, name)
-                # 修复图像处理
                 if buf:
-                    buf.seek(0)  # 重置缓冲区位置
+                    buf.seek(0)
                     img = Image.open(buf)
                     pred_text = f"{name} Result: {'1' if pred == 1 else '0'}"
-                    logger.info(f"{name} 预测结果: {pred}")
                     results.append((pred_text, img))
                 else:
-                    error_msg = f"{name} 预测失败: 未生成图像"
-                    logger.error(error_msg)
-                    results.append((error_msg, None))
-                    
+                    results.append((f"{name} 预测失败: 未生成图像", None))
             except Exception as e:
-                error_msg = f"{name} 预测过程中发生错误: {str(e)}"
-                logger.error(error_msg)
-                results.append((error_msg, None))
-        
-        return results[0][0], results[0][1], results[1][0], results[1][1], results[2][0], results[2][1]
-    
+                results.append((f"{name} 预测过程中发生错误: {str(e)}", None))
+        return (results[0][0], results[0][1],
+                results[1][0], results[1][1],
+                results[2][0], results[2][1])
     except Exception as e:
-        error_msg = f"预测过程中发生严重错误: {str(e)}"
-        logger.error(error_msg)
-        return error_msg, None, error_msg, None, error_msg, None
+        logger.error(f"预测过程中发生严重错误: {str(e)}")
+        return ("预测失败", None, "预测失败", None, "预测失败", None)
 
 # 原子和键类型选项
 ATOM_TYPES = ["C", "N", "O", "S", "P", "F", "Cl", "Br", "I", "H"]
@@ -206,318 +147,149 @@ BOND_TYPES = ["单键", "双键", "三键"]
 
 # 初始化分子结构
 def init_molecule():
-    return {
-        "atoms": [],
-        "bonds": []
-    }
+    return {"atoms": [], "bonds": []}
 
 # 添加原子
 def add_atom(molecule, atom_type):
-    molecule["atoms"].append({
-        "id": len(molecule["atoms"]),
-        "type": atom_type
-    })
+    molecule["atoms"].append({"id": len(molecule["atoms"]), "type": atom_type})
     return molecule
 
 # 添加键
 def add_bond(molecule, atom1_id, atom2_id, bond_type):
-    # 检查原子是否存在
     if atom1_id >= len(molecule["atoms"]) or atom2_id >= len(molecule["atoms"]):
         return molecule
-    
-    # 检查是否已存在键
     for bond in molecule["bonds"]:
-        if (bond["atom1"] == atom1_id and bond["atom2"] == atom2_id) or \
-           (bond["atom1"] == atom2_id and bond["atom2"] == atom1_id):
+        if {bond["atom1"], bond["atom2"]} == {atom1_id, atom2_id}:
             return molecule
-    
-    molecule["bonds"].append({
-        "atom1": atom1_id,
-        "atom2": atom2_id,
-        "type": bond_type
-    })
+    molecule["bonds"].append({"atom1": atom1_id, "atom2": atom2_id, "type": bond_type})
     return molecule
 
-# 从JSON结构生成SMILES
+# 从 JSON 构建 SMILES
 def generate_smiles(molecule_json):
     try:
-        # 创建空分子
         mol = Chem.RWMol()
-        
-        # 添加原子
         atom_map = {}
         for atom in molecule_json["atoms"]:
-            new_atom = Chem.Atom(atom["type"])
-            idx = mol.AddAtom(new_atom)
+            idx = mol.AddAtom(Chem.Atom(atom["type"]))
             atom_map[atom["id"]] = idx
-        
-        # 添加键
         for bond in molecule_json["bonds"]:
-            start_atom = atom_map[bond["atom1"]]
-            end_atom = atom_map[bond["atom2"]]
-            
-            # 确定键类型
-            bond_type_mapping = {
-                "单键": Chem.BondType.SINGLE,
-                "双键": Chem.BondType.DOUBLE,
-                "三键": Chem.BondType.TRIPLE
-            }
-            bond_type = bond_type_mapping.get(bond["type"], Chem.BondType.SINGLE)
-            
-            mol.AddBond(start_atom, end_atom, bond_type)
-        
-        # 清理分子
+            t = {"单键": Chem.BondType.SINGLE,
+                 "双键": Chem.BondType.DOUBLE,
+                 "三键": Chem.BondType.TRIPLE}[bond["type"]]
+            mol.AddBond(atom_map[bond["atom1"]], atom_map[bond["atom2"]], t)
         mol.UpdatePropertyCache()
         Chem.SanitizeMol(mol)
-        
-        # 生成SMILES
-        smiles = Chem.MolToSmiles(mol)
-        return smiles
+        return Chem.MolToSmiles(mol)
     except Exception as e:
-        logger.error(f"生成SMILES失败: {str(e)}")
+        logger.error(f"生成SMILES失败: {e}")
         return None
 
-# 可视化分子结构
+# 可视化分子
 def visualize_molecule(molecule_json):
-    try:
-        smiles = generate_smiles(molecule_json)
-        if not smiles:
-            return None
-        
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            return None
-        
-        # 生成2D坐标
-        AllChem.Compute2DCoords(mol)
-        
-        # 创建图像
-        img = Draw.MolToImage(mol, size=(300, 300))
-        img_buffer = BytesIO()
-        img.save(img_buffer, format="PNG")
-        img_buffer.seek(0)
-        return img_buffer
-    except Exception as e:
-        logger.error(f"可视化分子失败: {str(e)}")
+    smiles = generate_smiles(molecule_json)
+    if not smiles:
+        return None
+    mol = MolFromSmiles(smiles)
+    if mol is None:
         return None
+    AllChem.Compute2DCoords(mol)
+    img = Draw.MolToImage(mol, size=(300, 300))
+    buf = BytesIO()
+    img.save(buf, format="PNG")
+    buf.seek(0)
+    return buf
+
+# 更新下拉菜单选项（改为 gr.update）
+def update_atom_dropdowns(molecule):
+    choices = [f"{a['id']}: {a['type']}" for a in molecule["atoms"]]
+    return gr.update(choices=choices, value=None), gr.update(choices=choices, value=None)
+
+# 更新原子列表
+def update_atoms_list(molecule):
+    return [[a["id"], a["type"]] for a in molecule["atoms"]]
+
+# 更新键列表
+def update_bonds_list(molecule):
+    out = []
+    for b in molecule["bonds"]:
+        t1 = next(a["type"] for a in molecule["atoms"] if a["id"]==b["atom1"])
+        t2 = next(a["type"] for a in molecule["atoms"] if a["id"]==b["atom2"])
+        out.append([f"{b['atom1']}: {t1}", f"{b['atom2']}: {t2}", b["type"]])
+    return out
 
-# 使用TabbedInterface组织两种输入方式
 with gr.Blocks(title="CrystalGAT") as demo:
     gr.Markdown("# CrystalGAT")
-    gr.Markdown("输入SMILES字符串或构建分子结构，CrystalGAT将预测其弹性、塑性和脆性分类并可视化注意力权重")
+    gr.Markdown("输入SMILES或构建分子，预测弹性/塑性/脆性并可视化")
     
     with gr.Tab("SMILES输入"):
-        gr.Markdown("### 输入SMILES字符串")
-        smiles_input = gr.Textbox(
-            label="SMILES", 
-            placeholder="输入SMILES字符串，例如: CCO",
-            interactive=True
-        )
-        gr.Examples(
-            examples=[
-                ["CCO", "乙醇"],
-                ["C1=CC=CC=C1", "苯"],
-                ["CCOC(=O)C", "乙酸乙酯"]
-            ],
-            inputs=smiles_input,
-            label="示例"
-        )
+        smiles_input = gr.Textbox(label="SMILES", placeholder="例如: CCO")
         submit_btn1 = gr.Button("预测", variant="primary")
     
     with gr.Tab("构建分子"):
-        gr.Markdown("### 构建分子结构")
-        gr.Markdown("1. 添加原子：选择原子类型并点击'添加原子'按钮")
-        gr.Markdown("2. 添加键：选择两个原子和键类型，然后点击'添加键'按钮")
-        gr.Markdown("3. 完成后点击'生成分子'按钮预览并预测")
-        
-        # 初始化分子状态
         molecule_state = gr.State(init_molecule())
+        status_msg = gr.Textbox(label="状态", interactive=False, value="请添加原子开始")
         
         with gr.Row():
-            # 原子选择
-            with gr.Column():
-                gr.Markdown("### 添加原子")
-                atom_select = gr.Dropdown(
-                    label="选择原子类型",
-                    choices=ATOM_TYPES,
-                    value="C"
-                )
-                add_atom_btn = gr.Button("添加原子")
-                atoms_list = gr.Dataframe(
-                    label="原子列表",
-                    headers=["ID", "原子类型"],
-                    datatype=["number", "str"],
-                    interactive=False
-                )
-            
-            # 键选择
-            with gr.Column():
-                gr.Markdown("### 添加键")
-                # 创建原子选项列表
-                atom_options = gr.State([])
-                
-                atom1_select = gr.Dropdown(
-                    label="选择第一个原子",
-                    choices=[],
-                    interactive=True
-                )
-                atom2_select = gr.Dropdown(
-                    label="选择第二个原子",
-                    choices=[],
-                    interactive=True
-                )
-                bond_select = gr.Dropdown(
-                    label="选择键类型",
-                    choices=BOND_TYPES,
-                    value="单键"
-                )
-                add_bond_btn = gr.Button("添加键")
-                bonds_list = gr.Dataframe(
-                    label="键列表",
-                    headers=["原子1", "原子2", "键类型"],
-                    datatype=["str", "str", "str"],
-                    interactive=False
-                )
+            atom_select = gr.Dropdown(label="选择原子类型", choices=ATOM_TYPES, value="C")
+            add_atom_btn = gr.Button("添加原子")
+        atoms_list = gr.Dataframe(headers=["ID","原子类型"], datatype=["number","str"], interactive=False)
         
-        # 操作按钮
         with gr.Row():
-            clear_btn = gr.Button("清除所有")
-            generate_btn = gr.Button("生成分子")
+            atom1_select = gr.Dropdown(label="第一个原子", choices=[], value=None)
+            atom2_select = gr.Dropdown(label="第二个原子", choices=[], value=None)
+            bond_select = gr.Dropdown(label="键类型", choices=BOND_TYPES, value="单键")
+            add_bond_btn = gr.Button("添加键")
+        bonds_list = gr.Dataframe(headers=["原子1","原子2","键类型"], datatype=["str","str","str"], interactive=False)
         
-        # 分子预览
+        clear_btn = gr.Button("清除所有")
+        generate_btn = gr.Button("生成分子")
+        molecule_smiles = gr.Textbox(label="SMILES结果", interactive=False)
         molecule_img = gr.Image(label="分子预览", interactive=False)
-        molecule_smiles = gr.Textbox(label="生成的SMILES", interactive=False)
+        submit_btn2 = gr.Button("使用此分子预测", variant="primary")
+        
+        # 绑定事件
+        add_atom_btn.click(fn=lambda at, mol: add_atom(mol, at),
+                           inputs=[atom_select, molecule_state],
+                           outputs=molecule_state) \
+                      .then(fn=update_atoms_list, inputs=molecule_state, outputs=atoms_list) \
+                      .then(fn=update_atom_dropdowns, inputs=molecule_state, outputs=[atom1_select, atom2_select]) \
+                      .then(fn=lambda: "原子添加成功", outputs=status_msg)
+        
+        add_bond_btn.click(fn=lambda a1, a2, b, mol: (add_bond(mol,
+                                          int(a1.split(":")[0]),
+                                          int(a2.split(":")[0]), b)) if a1 and a2 else mol,
+                           inputs=[atom1_select, atom2_select, bond_select, molecule_state],
+                           outputs=molecule_state) \
+                      .then(fn=update_bonds_list, inputs=molecule_state, outputs=bonds_list) \
+                      .then(fn=lambda: "键添加/更新成功", outputs=status_msg)
         
-        # 预测按钮
-        submit_btn2 = gr.Button("使用此分子进行预测", variant="primary")
+        clear_btn.click(fn=init_molecule, outputs=molecule_state) \
+                 .then(fn=lambda: ([], []), outputs=[atoms_list, bonds_list]) \
+                 .then(fn=lambda: (gr.update(choices=[], value=None), gr.update(choices=[], value=None)),
+                       outputs=[atom1_select, atom2_select]) \
+                 .then(fn=lambda: "已清除所有", outputs=status_msg)
+        
+        generate_btn.click(fn=generate_smiles, inputs=molecule_state, outputs=molecule_smiles) \
+                    .then(fn=visualize_molecule, inputs=molecule_state, outputs=molecule_img) \
+                    .then(fn=lambda: "分子生成完成", outputs=status_msg)
     
-    # 输出区域 (两种输入方式共享)
-    gr.Markdown("## 预测结果")
+    # 输出区域
     with gr.Row():
         elastic_text = gr.Text(label="Elastic")
-        elastic_img = gr.Image(type="pil", label="Elastic attention visualization")
+        elastic_img = gr.Image(type="pil", label="Elastic 可视化")
     with gr.Row():
         plastic_text = gr.Text(label="Plastic")
-        plastic_img = gr.Image(type="pil", label="Plastic attention visualization")
+        plastic_img = gr.Image(type="pil", label="Plastic 可视化")
     with gr.Row():
         brittle_text = gr.Text(label="Brittle")
-        brittle_img = gr.Image(type="pil", label="Brittle attention visualization")
-    
-    # 更新原子选项列表
-    def update_atom_options(molecule):
-        # 创建原子选项列表
-        options = [f"{i}: {atom['type']}" for i, atom in enumerate(molecule["atoms"])]
-        return options
+        brittle_img = gr.Image(type="pil", label="Brittle 可视化")
     
-    # 更新原子列表显示
-    def update_atoms_list(molecule):
-        atoms_data = [[i, atom["type"]] for i, atom in enumerate(molecule["atoms"])]
-        return atoms_data
-    
-    # 更新键列表显示
-    def update_bonds_list(molecule):
-        bonds_data = []
-        for bond in molecule["bonds"]:
-            atom1_type = molecule["atoms"][bond["atom1"]]["type"]
-            atom2_type = molecule["atoms"][bond["atom2"]]["type"]
-            bonds_data.append([
-                f"{bond['atom1']}: {atom1_type}",
-                f"{bond['atom2']}: {atom2_type}",
-                bond["type"]
-            ])
-        return bonds_data
-    
-    # 事件处理 - 添加原子
-    add_atom_btn.click(
-        fn=lambda atom, mol: add_atom(mol, atom),
-        inputs=[atom_select, molecule_state],
-        outputs=molecule_state
-    ).then(
-        fn=update_atoms_list,
-        inputs=molecule_state,
-        outputs=atoms_list
-    ).then(
-        fn=update_atom_options,
-        inputs=molecule_state,
-        outputs=atom_options
-    ).then(
-        lambda options: [
-            gr.Dropdown.update(choices=options, value=options[0] if options else None),
-            gr.Dropdown.update(choices=options, value=options[0] if options else None)
-        ],
-        inputs=atom_options,
-        outputs=[atom1_select, atom2_select]
-    )
-    
-    # 事件处理 - 添加键
-    add_bond_btn.click(
-        fn=lambda atom1, atom2, bond, mol: add_bond(mol, int(atom1.split(":")[0]), int(atom2.split(":")[0]), bond),
-        inputs=[atom1_select, atom2_select, bond_select, molecule_state],
-        outputs=molecule_state
-    ).then(
-        fn=update_bonds_list,
-        inputs=molecule_state,
-        outputs=bonds_list
-    )
-    
-    # 事件处理 - 清除所有
-    clear_btn.click(
-        fn=init_molecule,
-        outputs=molecule_state
-    ).then(
-        fn=lambda: [],
-        outputs=atoms_list
-    ).then(
-        fn=lambda: [],
-        outputs=bonds_list
-    ).then(
-        fn=lambda: [],
-        outputs=atom_options
-    ).then(
-        lambda: [
-            gr.Dropdown.update(choices=[], value=None),
-            gr.Dropdown.update(choices=[], value=None)
-        ],
-        outputs=[atom1_select, atom2_select]
-    ).then(
-        fn=lambda: None,
-        outputs=molecule_img
-    ).then(
-        fn=lambda: "",
-        outputs=molecule_smiles
-    )
-    
-    # 事件处理 - 生成分子
-    generate_btn.click(
-        fn=generate_smiles,
-        inputs=molecule_state,
-        outputs=molecule_smiles
-    ).then(
-        fn=visualize_molecule,
-        inputs=molecule_state,
-        outputs=molecule_img
-    )
-    
-    # 设置交互
-    # SMILES输入路径
-    submit_btn1.click(
-        fn=predict_all,
-        inputs=smiles_input,
-        outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
-    )
-    
-    # 分子构建预测路径
-    submit_btn2.click(
-        fn=lambda smiles: predict_all(smiles) if smiles else ("请输入有效的SMILES", None, "", None, "", None),
-        inputs=molecule_smiles,
-        outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
-    )
+    submit_btn1.click(fn=predict_all,
+                      inputs=smiles_input,
+                      outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img])
+    submit_btn2.click(fn=lambda s: predict_all(s) if s else ("请输入SMILES", None, "", None, "", None),
+                      inputs=molecule_smiles,
+                      outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img])
 
 if __name__ == "__main__":
-    logger.info("启动Gradio应用...")
-    try:
-        demo.launch(server_name="0.0.0.0", server_port=7860)
-    except Exception as e:
-        logger.error(f"启动Gradio应用失败: {str(e)}")
-        logger.error(traceback.format_exc())
-        raise
+    demo.launch(server_name="0.0.0.0", server_port=7860)