Update app.py

app.py

@@ -20,8 +20,8 @@ from rdkit.Chem import AllChem
 from rdkit.Chem import MolFromSmiles
 from io import BytesIO
 import traceback
-import tempfile
 import base64
+import json
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
@@ -200,58 +200,236 @@ def predict_all(smiles):
         logger.error(error_msg)
         return error_msg, None, error_msg, None, error_msg, None
 
-# 新增函数：将绘制的分子结构转换为SMILES
-def draw_to_smiles(image_data):
-    """
-    将绘制的分子图像转换为SMILES字符串
-    :param image_data: 包含分子图像数据的base64字符串
-    :return: (SMILES字符串, 分子图像)
-    """
-    if not image_data:
-        return None, None
+# 原子和键类型选项
+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, x, y):
+    molecule["atoms"].append({
+        "id": len(molecule["atoms"]),
+        "type": atom_type,
+        "x": x,
+        "y": y
+    })
+    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):
+            return molecule
+    
+    molecule["bonds"].append({
+        "atom1": atom1_id,
+        "atom2": atom2_id,
+        "type": bond_type
+    })
+    return molecule
+
+# 从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)
+            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)
+        
+        # 清理分子
+        mol.UpdatePropertyCache()
+        Chem.SanitizeMol(mol)
         
+        # 生成SMILES
+        smiles = Chem.MolToSmiles(mol)
+        return smiles
+    except Exception as e:
+        logger.error(f"生成SMILES失败: {str(e)}")
+        return None
+
+# 可视化分子结构
+def visualize_molecule(molecule_json):
     try:
-        # 将base64图像数据转换为PIL图像
-        image_bytes = base64.b64decode(image_data.split(",")[1])
-        image = Image.open(BytesIO(image_bytes))
+        smiles = generate_smiles(molecule_json)
+        if not smiles:
+            return None
         
-        # 保存临时文件
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            image.save(tmp.name)
-            tmp_path = tmp.name
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
         
-        # 使用RDKit从图像读取分子
-        mol = MolFromSmiles("")  # 创建一个空分子
-        # 注意：实际应用中这里应该使用OCR或图像识别技术
-        # 这里简化处理，提示用户输入SMILES
+        # 生成2D坐标
+        AllChem.Compute2DCoords(mol)
         
-        # 创建分子图像
+        # 创建图像
+        img = Draw.MolToImage(mol, size=(300, 300))
         img_buffer = BytesIO()
-        image.save(img_buffer, format="PNG")
+        img.save(img_buffer, format="PNG")
         img_buffer.seek(0)
+        return img_buffer
+    except Exception as e:
+        logger.error(f"可视化分子失败: {str(e)}")
+        return None
+
+# 分子构建界面
+def build_molecule_interface():
+    # 初始化分子
+    molecule = init_molecule()
+    
+    with gr.Blocks() as interface:
+        gr.Markdown("### 构建分子")
+        
+        # 原子选择
+        with gr.Row():
+            atom_select = gr.Dropdown(
+                label="选择原子类型",
+                choices=ATOM_TYPES,
+                value="C"
+            )
+            bond_select = gr.Dropdown(
+                label="选择键类型",
+                choices=BOND_TYPES,
+                value="单键"
+            )
+        
+        # 画布
+        canvas = gr.Image(
+            label="分子结构",
+            interactive=True,
+            tool="select",
+            height=400
+        )
+        
+        # 状态显示
+        status = gr.Textbox(label="状态", interactive=False)
+        
+        # 操作按钮
+        with gr.Row():
+            add_atom_btn = gr.Button("添加原子")
+            add_bond_btn = gr.Button("添加键")
+            clear_btn = gr.Button("清除")
+            generate_btn = gr.Button("生成分子")
+        
+        # 分子预览
+        molecule_img = gr.Image(label="分子预览", interactive=False)
+        molecule_smiles = gr.Textbox(label="生成的SMILES", interactive=False)
+        
+        # 存储分子结构的隐藏状态
+        molecule_state = gr.State(molecule)
         
-        # 提示用户输入SMILES
-        return None, img_buffer
+        # 原子位置存储
+        last_click_pos = gr.State((0, 0))
+        
+        # 事件处理
+        canvas.select(
+            fn=lambda evt: {"x": evt.index[0], "y": evt.index[1]},
+            outputs=last_click_pos
+        )
+        
+        add_atom_btn.click(
+            fn=lambda atom, pos, mol: add_atom(mol, atom, pos[0], pos[1]),
+            inputs=[atom_select, last_click_pos, molecule_state],
+            outputs=molecule_state
+        ).then(
+            fn=lambda mol: f"添加原子成功! 当前原子数: {len(mol['atoms'])}, 键数: {len(mol['bonds'])}",
+            inputs=molecule_state,
+            outputs=status
+        )
+        
+        # 添加键需要选择两个原子
+        atom1_state = gr.State(-1)
+        atom2_state = gr.State(-1)
+        
+        canvas.select(
+            fn=lambda evt, mol: {"atom_id": find_nearest_atom(mol, evt.index[0], evt.index[1])},
+            inputs=molecule_state,
+            outputs=atom1_state
+        )
+        
+        canvas.select(
+            fn=lambda evt, mol: {"atom_id": find_nearest_atom(mol, evt.index[0], evt.index[1])},
+            inputs=molecule_state,
+            outputs=atom2_state
+        )
+        
+        add_bond_btn.click(
+            fn=lambda bond, atom1, atom2, mol: add_bond(mol, atom1, atom2, bond),
+            inputs=[bond_select, atom1_state, atom2_state, molecule_state],
+            outputs=molecule_state
+        ).then(
+            fn=lambda mol: f"添加键成功! 当前原子数: {len(mol['atoms'])}, 键数: {len(mol['bonds'])}",
+            inputs=molecule_state,
+            outputs=status
+        )
+        
+        clear_btn.click(
+            fn=init_molecule,
+            outputs=molecule_state
+        ).then(
+            fn=lambda: "已清除所有原子和键",
+            outputs=status
+        )
+        
+        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 mol: f"分子生成完成! 原子数: {len(mol['atoms'])}, 键数: {len(mol['bonds'])}",
+            inputs=molecule_state,
+            outputs=status
+        )
     
-    except Exception as e:
-        error_msg = f"分子转换失败: {str(e)}"
-        logger.error(error_msg)
-        return None, None
+    return interface
 
-# 新增函数：处理绘制输入的预测流程
-def predict_from_drawing(smiles):
-    """
-    处理用户绘制的分子结构
-    :param smiles: 用户输入的SMILES字符串
-    :return: 预测结果 (6个输出组件)
-    """
-    # 使用预测函数
-    return predict_all(smiles)
+# 查找最近的原子
+def find_nearest_atom(molecule, x, y, radius=20):
+    for atom in molecule["atoms"]:
+        dist = ((atom["x"] - x)**2 + (atom["y"] - y)**2)**0.5
+        if dist <= radius:
+            return atom["id"]
+    return -1
 
 # 使用TabbedInterface组织两种输入方式
 with gr.Blocks(title="CrystalGAT") as demo:
     gr.Markdown("# CrystalGAT")
-    gr.Markdown("输入SMILES字符串或绘制分子结构，CrystalGAT将预测其弹性、塑性和脆性分类并可视化注意力权重")
+    gr.Markdown("输入SMILES字符串或构建分子结构，CrystalGAT将预测其弹性、塑性和脆性分类并可视化注意力权重")
     
     with gr.Tab("SMILES输入"):
         gr.Markdown("### 输入SMILES字符串")
@@ -271,27 +449,19 @@ with gr.Blocks(title="CrystalGAT") as demo:
         )
         submit_btn1 = gr.Button("预测", variant="primary")
     
-    with gr.Tab("绘制分子"):
-        gr.Markdown("### 绘制分子结构")
-        gr.Markdown("使用下方的绘图工具绘制分子结构，然后在SMILES框中输入对应的SMILES字符串")
-        
-        # 使用标准绘图组件
-        drawing_input = gr.ImageEditor(
-            label="绘制分子结构",
-            type="pil",
-            interactive=True,
-            height=300
-        )
+    with gr.Tab("构建分子"):
+        gr.Markdown("### 构建分子结构")
+        gr.Markdown("1. 选择原子类型和键类型")
+        gr.Markdown("2. 点击'添加原子'按钮后在画布上点击放置原子")
+        gr.Markdown("3. 点击'添加键'按钮后依次点击两个原子连接它们")
+        gr.Markdown("4. 完成后点击'生成分子'按钮")
         
-        gr.Markdown("### 输入绘制的分子对应的SMILES")
-        drawing_smiles = gr.Textbox(
-            label="SMILES", 
-            placeholder="输入绘制的分子对应的SMILES字符串",
-            interactive=True
-        )
+        # 分子构建界面
+        build_interface = build_molecule_interface()
         
-        drawing_display = gr.Image(label="分子结构预览", interactive=False)
-        submit_btn2 = gr.Button("预测", variant="primary")
+        # 预测按钮
+        submit_btn2 = gr.Button("使用此分子进行预测", variant="primary")
+        build_smiles = gr.Textbox(visible=False)  # 隐藏的SMILES存储
     
     # 输出区域 (两种输入方式共享)
     gr.Markdown("## 预测结果")
@@ -313,16 +483,10 @@ with gr.Blocks(title="CrystalGAT") as demo:
         outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
     )
     
-    # 绘图输入路径
-    drawing_input.change(
-        fn=draw_to_smiles,
-        inputs=drawing_input,
-        outputs=[drawing_smiles, drawing_display]
-    )
-    
+    # 分子构建预测路径
     submit_btn2.click(
-        fn=predict_from_drawing,
-        inputs=drawing_smiles,
+        fn=lambda smiles: predict_all(smiles) if smiles else ("请输入有效的SMILES", None, "", None, "", None),
+        inputs=build_smiles,
         outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
     )