Update app.py

app.py

@@ -1,224 +1,303 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Tue Jul  8 15:00:01 2025
-
-@author: User
-"""
-import os
-import zipfile
-import torch
-import numpy as np
-import logging
-from PIL import Image
-import gradio as gr
-from torch_geometric.data import Data as PyGData
-import matplotlib
-matplotlib.use('Agg')  # 修复后台线程问题
-
-# 配置日志
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# GPU内存优化
-os.environ["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:
-        with zipfile.ZipFile("models.zip", 'r') as zip_ref:
-            zip_ref.extractall(".")
-        logger.info("模型文件解压完成!")
-    except Exception as e:
-        logger.error(f"解压模型文件失败: {str(e)}")
-        raise
-
-# 导入模型工具
-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)}")
-    raise
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-logger.info(f"使用设备: {device}")
-if torch.cuda.is_available():
-    logger.info(f"GPU 信息: {torch.cuda.get_device_name(0)}")
-
-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)
-    }
-    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
-                    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: 添加安全全局变量
-                try:
-                    # 导入必要的模块
-                    import numpy as np
-                    import torch.serialization
-                    
-                    # 添加安全全局变量
-                    torch.serialization.add_safe_globals([getattr(np.core.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("使用整个状态字典加载模型")
-            
-            # 加载模型参数
-            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
-
-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)
-        )
-
-        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'}"
-                    logger.info(f"{name} 预测结果: {pred}")
-                    results.append((pred_text, img))
-                else:
-                    error_msg = f"{name} 预测失败: 未生成图像"
-                    logger.error(error_msg)
-                    results.append((error_msg, 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]
-    
-    except Exception as e:
-        error_msg = f"预测过程中发生严重错误: {str(e)}"
-        logger.error(error_msg)
-        return error_msg, None, error_msg, None, error_msg, None
-
-# 修复输出组件
-outputs = [
-    gr.Text(label="Elastic"), 
-    gr.Image(type="pil", label="Elastic attention visualization"),
-    gr.Text(label="Plastic"), 
-    gr.Image(type="pil", label="Plastic attention visualization"),
-    gr.Text(label="Brittle"), 
-    gr.Image(type="pil", label="Brittle attention visualization")
-]
-
-demo = gr.Interface(
-    fn=predict_all,
-    inputs=gr.Textbox(label="SMILES", placeholder="Enter a SMILES string, for example: CCO"),
-    outputs=outputs,
-    title="CrystalGAT",
-    description="Input the SMILES string of a molecule, and CrystalGAT will predict its elasticity, plasticity and brittleness classification and visualize the attention weights",
-    examples=[
-        ["CCO", "乙醇"],
-        ["C1=CC=CC=C1", "苯"],
-        ["CCOC(=O)C", "乙酸乙酯"]
-    ]
-)
-
-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)}")
-        raise
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jul  8 15:00:01 2025
+@author: User
+"""
+import os
+import zipfile
+import torch
+import numpy as np
+import logging
+from PIL import Image
+import gradio as gr
+from torch_geometric.data import Data as PyGData
+import matplotlib
+matplotlib.use('Agg')  # 修复后台线程问题
+
+# 新增导入
+from rdkit import Chem
+from rdkit.Chem import AllChem, Draw
+from rdkit.Chem.Draw import rdMolDraw2D
+import tempfile
+
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# GPU内存优化
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+    logger.info("设置GPU内存优化参数: PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128")
+else:
+    logger.info("使用CPU运行")
+
+# 解压模型文件
+MODEL_FILES = ["best_model-E-500-68.pth", "best_model-P-5000-180.pth", "best_model-B-6000-185.pth"]
+if not all(os.path.exists(f) for f in MODEL_FILES):
+    logger.info("开始解压模型文件...")
+    try:
+        with zipfile.ZipFile("models.zip", 'r') as zip_ref:
+            zip_ref.extractall(".")
+        logger.info("模型文件解压完成!")
+    except Exception as e:
+        logger.error(f"解压模型文件失败: {str(e)}")
+        raise
+
+# 导入模型工具
+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)}")
+    raise
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+logger.info(f"使用设备: {device}")
+if torch.cuda.is_available():
+    logger.info(f"GPU 信息: {torch.cuda.get_device_name(0)}")
+
+def load_models():
+    model_info = {
+        "Elastic": ("best_model-E-500-68.pth", 2),
+        "Plastic": ("best_model-P-5000-180.pth", 2),
+        "Brittle": ("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}")
+            raise FileNotFoundError(f"模型文件 {pth_path} 不存在")
+        
+        try:
+            model = EnhancedGAT(input_dim=12, hidden_dim=512, output_dim=output_dim, num_heads=8)
+            
+            # 简化模型加载
+            state_dict = torch.load(pth_path, map_location=device)
+            if "model_state_dict" in state_dict:
+                model.load_state_dict(state_dict["model_state_dict"])
+            else:
+                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)}")
+    # 创建虚拟模型保持应用运行
+    dummy_model = EnhancedGAT(input_dim=12, hidden_dim=512, output_dim=2, num_heads=8).eval()
+    models = {name: dummy_model for name in ["Elastic", "Plastic", "Brittle"]}
+    logger.warning("使用虚拟模型继续运行，功能受限")
+
+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
+        )
+
+        results = []
+        # 对每个模型进行预测
+        for name in ["Elastic", "Plastic", "Brittle"]:
+            logger.info(f"使用 {name} 模型进行预测...")
+            try:
+                buf, pred = visualize_single_molecule(models[name], data, device, name)
+                if buf:
+                    img = Image.open(buf)
+                    pred_text = f"{name}: {'Positive' if pred == 1 else 'Negative'}"
+                    results.append((pred_text, img))
+                else:
+                    error_msg = f"{name} 预测失败: 未生成图像"
+                    logger.error(error_msg)
+                    results.append((error_msg, 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]
+    except Exception as e:
+        error_msg = f"预测过程错误: {str(e)}"
+        logger.error(error_msg)
+        return error_msg, None, error_msg, None, error_msg, None
+
+# ===== 新增分子结构绘制功能 =====
+def draw_molecule(smiles=None):
+    """根据SMILES生成分子结构图像"""
+    try:
+        if not smiles:
+            return None
+        
+        mol = Chem.MolFromSmiles(smiles)
+        if not mol:
+            return None
+            
+        # 生成2D分子图像
+        drawer = rdMolDraw2D.MolDraw2DCairo(400, 300)
+        drawer.DrawMolecule(mol)
+        drawer.FinishDrawing()
+        
+        # 转换为PIL图像
+        img_data = drawer.GetDrawingText()
+        return Image.open(tempfile.NamedTemporaryFile(delete=False, suffix='.png').write(img_data))
+    except Exception as e:
+        logger.error(f"分子绘制失败: {str(e)}")
+        return None
+
+def predict_from_structure(mol_dict):
+    """从绘制的分子结构预测"""
+    try:
+        if not mol_dict or not mol_dict['atoms']:
+            return "请绘制分子结构", None, "请绘制分子结构", None, "请绘制分子结构", None
+        
+        # 转换绘制结构为SMILES
+        mol = Chem.RWMol()
+        atom_map = {}
+        
+        # 添加原子
+        for atom in mol_dict['atoms']:
+            new_atom = Chem.Atom(atom['atom_symbol'])
+            idx = mol.AddAtom(new_atom)
+            atom_map[atom['atom_index']] = idx
+        
+        # 添加键
+        for bond in mol_dict['bonds']:
+            start = atom_map[bond['start_atom']]
+            end = atom_map[bond['end_atom']]
+            bond_type = Chem.BondType.values[bond['bond_type'] - 1]  # 转换键类型
+            mol.AddBond(start, end, bond_type)
+        
+        # 获取SMILES
+        smiles = Chem.MolToSmiles(mol)
+        logger.info(f"转换的SMILES: {smiles}")
+        
+        # 进行预测
+        return predict_all(smiles)
+    except Exception as e:
+        error_msg = f"结构转换错误: {str(e)}"
+        logger.error(error_msg)
+        return error_msg, None, error_msg, None, error_msg, None
+
+# ===== 创建多选项卡界面 =====
+with gr.Blocks(title="CrystalGAT") as demo:
+    gr.Markdown("# CrystalGAT分子性质预测")
+    gr.Markdown("输入SMILES或绘制分子结构，预测弹性、塑性和脆性分类并可视化注意力权重")
+    
+    with gr.Tab("SMILES输入"):
+        with gr.Row():
+            with gr.Column():
+                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
+                )
+                submit_btn = gr.Button("预测")
+            with gr.Column():
+                molecule_img = gr.Image(label="分子结构", interactive=False)
+                smiles_input.change(
+                    fn=draw_molecule, 
+                    inputs=smiles_input, 
+                    outputs=molecule_img
+                )
+        
+        with gr.Row():
+            with gr.Column():
+                elastic_text = gr.Text(label="弹性")
+                elastic_img = gr.Image(label="注意力可视化")
+            with gr.Column():
+                plastic_text = gr.Text(label="塑性")
+                plastic_img = gr.Image(label="注意力可视化")
+            with gr.Column():
+                brittle_text = gr.Text(label="脆性")
+                brittle_img = gr.Image(label="注意力可视化")
+        
+        submit_btn.click(
+            fn=predict_all,
+            inputs=smiles_input,
+            outputs=[
+                elastic_text, elastic_img, 
+                plastic_text, plastic_img,
+                brittle_text, brittle_img
+            ]
+        )
+    
+    with gr.Tab("绘制分子结构"):
+        with gr.Row():
+            with gr.Column():
+                molecule_editor = gr.Molecule(
+                    label="绘制分子结构",
+                    type="sketch",
+                    interactive=True
+                )
+                draw_submit = gr.Button("预测")
+            with gr.Column():
+                gr.Markdown("### 绘制说明")
+                gr.Markdown("1. 从右侧选择原子工具<br>2. 在画布上点击添加原子<br>3. 选择键工具连接原子<br>4. 点击预测按钮进行分析")
+        
+        with gr.Row():
+            with gr.Column():
+                draw_elastic_text = gr.Text(label="弹性")
+                draw_elastic_img = gr.Image(label="注意力可视化")
+            with gr.Column():
+                draw_plastic_text = gr.Text(label="塑性")
+                draw_plastic_img = gr.Image(label="注意力可视化")
+            with gr.Column():
+                draw_brittle_text = gr.Text(label="脆性")
+                draw_brittle_img = gr.Image(label="注意力可视化")
+        
+        draw_submit.click(
+            fn=predict_from_structure,
+            inputs=molecule_editor,
+            outputs=[
+                draw_elastic_text, draw_elastic_img, 
+                draw_plastic_text, draw_plastic_img,
+                draw_brittle_text, draw_brittle_img
+            ]
+        )
+
+if __name__ == "__main__":
+    logger.info("启动Gradio应用...")
+    try:
+        demo.launch(
+            server_name="0.0.0.0", 
+            server_port=7860,
+            share=False,
+            show_error=True
+        )
+    except Exception as e:
+        logger.error(f"启动Gradio应用失败: {str(e)}")
+        raise