Update app.py

app.py

@@ -17,18 +17,19 @@ matplotlib.use('Agg')  # 修复后台线程问题
 from rdkit import Chem
 from rdkit.Chem import Draw, AllChem, MolFromSmiles
 from io import BytesIO
-import traceback
 
 # 配置日志
-logging.basicConfig(level=logging.INFO,
-                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
-# GPU内存优化
+# GPU 内存优化
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
-logger.info("设置GPU内存优化参数: PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128")
+logger.info("设置 GPU 内存优化参数: PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128")
 
-# 解压模型文件
+# 解压模型文件（如果尚未解压）
 if not os.path.exists("best_model-B-6000-185.pth"):
     logger.info("开始解压模型文件...")
     try:
@@ -36,7 +37,7 @@ if not os.path.exists("best_model-B-6000-185.pth"):
             zip_ref.extractall(".")
         logger.info("模型文件解压完成!")
     except Exception as e:
-        logger.error(f"解压模型文件失败: {str(e)}")
+        logger.error(f"解压模型文件失败: {e}")
         raise
 
 # 导入模型工具
@@ -44,7 +45,7 @@ try:
     from model_utils import EnhancedGAT, smiles_to_graph, visualize_single_molecule
     logger.info("成功导入 model_utils 模块")
 except ImportError as e:
-    logger.error(f"导入 model_utils 失败: {str(e)}")
+    logger.error(f"导入 model_utils 失败: {e}")
     raise
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -56,263 +57,199 @@ def load_models():
     model_info = {
         "Elastic": ("models/best_model-E-500-68.pth", 2),
         "Plastic": ("models/best_model-P-5000-180.pth", 2),
-        "Brittle": ("models/best_model-B-6000-185.pth", 2)
+        "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(),
-                pth_path.upper(),
-                pth_path.replace("-", "_"),
-                pth_path.replace("_", "-")
-            ]
-            found = False
-            for file in possible_files:
-                if os.path.exists(file):
-                    logger.warning(f"使用替代文件: {file}")
-                    pth_path = file
-                    found = True
+            # 尝试其他变体
+            for alt in [pth_path.lower(), pth_path.upper(), pth_path.replace("-", "_"), pth_path.replace("_", "-")]:
+                if os.path.exists(alt):
+                    logger.warning(f"使用替代模型文件: {alt}")
+                    pth_path = alt
                     break
-            if not found:
+            else:
                 raise FileNotFoundError(f"模型文件 {pth_path} 不存在")
+        model = EnhancedGAT(input_dim=12, hidden_dim=512, output_dim=output_dim, num_heads=8)
         try:
-            model = EnhancedGAT(input_dim=12, hidden_dim=512,
-                                output_dim=output_dim, num_heads=8)
-            logger.info(f"加载模型权重: {pth_path}")
-            try:
-                state = torch.load(pth_path, map_location=device, weights_only=False)
-            except Exception as e:
-                logger.warning(f"weights_only=False 加载失败: {str(e)}")
-                try:
-                    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)
-                except:
-                    state = torch.load(pth_path, map_location=device)
-            if "model_state_dict" in state:
-                state_dict = state["model_state_dict"]
-            else:
-                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
+            state = torch.load(pth_path, map_location=device, weights_only=False)
+        except Exception:
+            # 备用加载方式
+            state = torch.load(pth_path, map_location=device)
+        # 支持两种格式
+        state_dict = state.get("model_state_dict", state)
+        model.load_state_dict(state_dict)
+        model.eval().to(device)
+        models[name] = model
+        logger.info(f"{name} 模型加载成功")
     return models
 
-logger.info("开始加载所有模型...")
 models = load_models()
-logger.info("所有模型加载完成!")
 
-def predict_all(smiles):
-    logger.info(f"收到预测请求: SMILES = {smiles}")
-    try:
-        atom_features, (rows, cols, edge_attr), mol = smiles_to_graph(smiles)
-        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))
-        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)
-                    img = Image.open(buf)
-                    pred_text = f"{name} Result: {'1' if pred == 1 else '0'}"
-                    results.append((pred_text, img))
-                else:
-                    results.append((f"{name} 预测失败: 未生成图像", None))
-            except Exception as e:
-                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:
-        logger.error(f"预测过程中发生严重错误: {str(e)}")
-        return ("预测失败", None, "预测失败", None, "预测失败", None)
+def predict_all(smiles: str):
+    """对 Elastic, Plastic, Brittle 三个模型做预测，返回文本与 PIL 图像对象。"""
+    atom_features, (rows, cols, edge_attr), mol = smiles_to_graph(smiles)
+    x = torch.tensor(atom_features, dtype=torch.float)
+    edge_index = torch.tensor(np.vstack((rows, cols)), 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),
+    )
+
+    outputs = []
+    for name in ["Elastic", "Plastic", "Brittle"]:
+        buf, pred = visualize_single_molecule(models[name], data, device, name)
+        if buf is not None:
+            buf.seek(0)
+            img = Image.open(buf)
+            text = f"{name} Result: {int(pred)}"
+            outputs.append((text, img))
+        else:
+            outputs.append((f"{name} 预测失败", None))
+    # 拆包为 6 个输出
+    return (*outputs[0], *outputs[1], *outputs[2])
+
+# 分子构建相关函数
 
-# 原子和键类型选项
 ATOM_TYPES = ["C", "N", "O", "S", "P", "F", "Cl", "Br", "I", "H"]
 BOND_TYPES = ["单键", "双键", "三键"]
 
-# 初始化分子结构
 def init_molecule():
     return {"atoms": [], "bonds": []}
 
-# 添加原子
-def add_atom(molecule, 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"], bond["atom2"]} == {atom1_id, atom2_id}:
-            return molecule
-    molecule["bonds"].append({"atom1": atom1_id, "atom2": atom2_id, "type": bond_type})
-    return molecule
-
-# 从 JSON 构建 SMILES
-def generate_smiles(molecule_json):
+def add_atom(mol_json, atom_type):
+    mol_json["atoms"].append({"id": len(mol_json["atoms"]), "type": atom_type})
+    return mol_json
+
+def add_bond(mol_json, atom1_sel, atom2_sel, bond_type):
+    if not atom1_sel or not atom2_sel:
+        return mol_json
+    a1 = int(atom1_sel.split(":")[0])
+    a2 = int(atom2_sel.split(":")[0])
+    # 避免重复
+    for b in mol_json["bonds"]:
+        if set([b["atom1"], b["atom2"]]) == set([a1, a2]):
+            return mol_json
+    mol_json["bonds"].append({"atom1": a1, "atom2": a2, "type": bond_type})
+    return mol_json
+
+def generate_smiles(mol_json):
     try:
         mol = Chem.RWMol()
-        atom_map = {}
-        for atom in molecule_json["atoms"]:
+        id_map = {}
+        for atom in mol_json["atoms"]:
             idx = mol.AddAtom(Chem.Atom(atom["type"]))
-            atom_map[atom["id"]] = idx
-        for bond in molecule_json["bonds"]:
-            t = {"单键": Chem.BondType.SINGLE,
-                 "双键": Chem.BondType.DOUBLE,
-                 "三键": Chem.BondType.TRIPLE}[bond["type"]]
-            mol.AddBond(atom_map[bond["atom1"]], atom_map[bond["atom2"]], t)
+            id_map[atom["id"]] = idx
+        for bond in mol_json["bonds"]:
+            bt = {"单键": Chem.BondType.SINGLE,
+                  "双键": Chem.BondType.DOUBLE,
+                  "三键": Chem.BondType.TRIPLE}[bond["type"]]
+            mol.AddBond(id_map[bond["atom1"]], id_map[bond["atom2"]], bt)
         mol.UpdatePropertyCache()
         Chem.SanitizeMol(mol)
         return Chem.MolToSmiles(mol)
     except Exception as e:
-        logger.error(f"生成SMILES失败: {e}")
-        return None
+        logger.error(f"生成 SMILES 失败: {e}")
+        return ""
 
-# 可视化分子
-def visualize_molecule(molecule_json):
-    try:
-        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
-    except Exception as e:
-        logger.error(f"可视化分子失败: {str(e)}")
+def visualize_molecule(mol_json):
+    """直接返回 PIL.Image.Image 对象，或 None"""
+    smiles = generate_smiles(mol_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))
+    return img
 
-# 更新下拉菜单
-def update_atom_dropdowns(molecule):
-    choices = [f"{a['id']}: {a['type']}" for a in molecule["atoms"]]
+def update_atom_dropdowns(mol_json):
+    choices = [f"{a['id']}: {a['type']}" for a in mol_json["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_atoms_list(mol_json):
+    return [[a["id"], a["type"]] for a in mol_json["atoms"]]
 
-def update_bonds_list(molecule):
+def update_bonds_list(mol_json):
     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"])
+    for b in mol_json["bonds"]:
+        t1 = next(a["type"] for a in mol_json["atoms"] if a["id"] == b["atom1"])
+        t2 = next(a["type"] for a in mol_json["atoms"] if a["id"] == b["atom2"])
         out.append([f"{b['atom1']}: {t1}", f"{b['atom2']}: {t2}", b["type"]])
     return out
 
 with gr.Blocks(title="CrystalGAT") as demo:
     gr.Markdown("# CrystalGAT")
-    gr.Markdown("输入SMILES或构建分子，预测弹性/塑性/脆性并可视化")
+    gr.Markdown("输入 SMILES 或 构建分子，预测弹性/塑性/脆性 并可视化注意力权重")
 
-    # SMILES 输入
-    with gr.Tab("SMILES输入"):
+    with gr.Tab("SMILES 输入"):
         smiles_input = gr.Textbox(label="SMILES", placeholder="例如: CCO")
-        submit_btn1 = gr.Button("预测", variant="primary")
+        predict_btn1 = gr.Button("预测", variant="primary")
 
-    # 构建分子
     with gr.Tab("构建分子"):
-        molecule_state = gr.State(init_molecule())
-        status_msg = gr.Textbox(label="状态", interactive=False, value="请添加原子开始")
+        state = gr.State(init_molecule())
+        status = gr.Textbox(label="状态", interactive=False, value="请添加原子开始")
 
         with gr.Row():
-            atom_select = gr.Dropdown(label="选择原子类型", choices=ATOM_TYPES, value="C")
+            atom_type = gr.Dropdown(label="选择原子类型", choices=ATOM_TYPES, value="C")
             add_atom_btn = gr.Button("添加原子")
-        atoms_list = gr.Dataframe(headers=["ID","原子类型"], datatype=["number","str"], interactive=False)
+        atom_table = gr.Dataframe(headers=["ID", "原子类型"], datatype=["number","str"], interactive=False)
 
         with gr.Row():
-            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="单键")
+            atom1 = gr.Dropdown(label="第一个原子", choices=[], value=None)
+            atom2 = gr.Dropdown(label="第二个原子", choices=[], value=None)
+            bond_type = 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)
+        bond_table = 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)
-        # 关键点：将 type="pil"，才能显示 PIL 图像
-        molecule_img = gr.Image(type="pil", label="分子预览", 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]) if a1 else -1,
-                                int(a2.split(":")[0]) if a2 else -1,
-                                b),
-                           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)
-
-        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)
-
-    # 预测结果展示
+        gen_btn = gr.Button("生成分子")
+        smiles_out = gr.Textbox(label="SMILES 结果", interactive=False)
+        mol_img = gr.Image(type="pil", label="分子预览")
+        predict_btn2 = gr.Button("使用此分子预测", variant="primary")
+
+        add_atom_btn.click(fn=add_atom, inputs=[state, atom_type], outputs=state) \
+            .then(fn=update_atoms_list, inputs=state, outputs=atom_table) \
+            .then(fn=update_atom_dropdowns, inputs=state, outputs=[atom1, atom2]) \
+            .then(fn=lambda: "原子添加成功", outputs=status)
+
+        add_bond_btn.click(fn=add_bond, inputs=[state, atom1, atom2, bond_type], outputs=state) \
+            .then(fn=update_bonds_list, inputs=state, outputs=bond_table) \
+            .then(fn=lambda: "键添加/更新成功", outputs=status)
+
+        clear_btn.click(fn=lambda: init_molecule(), outputs=state) \
+            .then(fn=lambda: ([], []), outputs=[atom_table, bond_table]) \
+            .then(fn=lambda: (gr.update(choices=[], value=None),
+                              gr.update(choices=[], value=None)),
+                  outputs=[atom1, atom2]) \
+            .then(fn=lambda: "已清除所有", outputs=status)
+
+        gen_btn.click(fn=generate_smiles, inputs=state, outputs=smiles_out) \
+               .then(fn=visualize_molecule, inputs=state, outputs=mol_img) \
+               .then(fn=lambda: "分子生成完成", outputs=status)
+
     with gr.Row():
-        elastic_text = gr.Text(label="Elastic")
-        elastic_img = gr.Image(type="pil", label="Elastic 可视化")
+        e_txt = gr.Text(label="Elastic")
+        e_img = gr.Image(type="pil", label="Elastic 可视化")
     with gr.Row():
-        plastic_text = gr.Text(label="Plastic")
-        plastic_img = gr.Image(type="pil", label="Plastic 可视化")
+        p_txt = gr.Text(label="Plastic")
+        p_img = gr.Image(type="pil", label="Plastic 可视化")
     with gr.Row():
-        brittle_text = gr.Text(label="Brittle")
-        brittle_img = gr.Image(type="pil", label="Brittle 可视化")
-
-    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])
+        b_txt = gr.Text(label="Brittle")
+        b_img = gr.Image(type="pil", label="Brittle 可视化")
+
+    predict_btn1.click(fn=predict_all,
+                       inputs=smiles_input,
+                       outputs=[e_txt, e_img, p_txt, p_img, b_txt, b_img])
+    predict_btn2.click(fn=lambda s: predict_all(s) if s else ("请输入SMILES", None, "", None, "", None),
+                       inputs=smiles_out,
+                       outputs=[e_txt, e_img, p_txt, p_img, b_txt, b_img])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)