Update app.py

app.py

@@ -20,6 +20,8 @@ from rdkit.Chem import AllChem
 from rdkit.Chem import 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')
@@ -331,9 +333,6 @@ with gr.Blocks(title="CrystalGAT") as demo:
         # 初始化分子状态
         molecule_state = gr.State(init_molecule())
         
-        # 状态消息
-        status_msg = gr.Textbox(label="状态", interactive=False, value="请添加原子开始构建分子")
-        
         with gr.Row():
             # 原子选择
             with gr.Column():
@@ -354,6 +353,9 @@ with gr.Blocks(title="CrystalGAT") as demo:
             # 键选择
             with gr.Column():
                 gr.Markdown("### 添加键")
+                # 创建原子选项列表
+                atom_options = gr.State([])
+                
                 atom1_select = gr.Dropdown(
                     label="选择第一个原子",
                     choices=[],
@@ -373,7 +375,7 @@ with gr.Blocks(title="CrystalGAT") as demo:
                 bonds_list = gr.Dataframe(
                     label="键列表",
                     headers=["原子1", "原子2", "键类型"],
-                    datatype=["number", "number", "str"],
+                    datatype=["str", "str", "str"],
                     interactive=False
                 )
         
@@ -401,32 +403,31 @@ with gr.Blocks(title="CrystalGAT") as demo:
         brittle_text = gr.Text(label="Brittle")
         brittle_img = gr.Image(type="pil", label="Brittle attention visualization")
     
-    # 更新原子选择下拉菜单
-    def update_atom_dropdowns(molecule):
-        atom_choices = [f"{atom['id']}: {atom['type']}" for atom in molecule["atoms"]]
-        return gr.Dropdown.update(choices=atom_choices), gr.Dropdown.update(choices=atom_choices)
+    # 更新原子选项列表
+    def update_atom_options(molecule):
+        # 创建原子选项列表
+        options = [f"{i}: {atom['type']}" for i, atom in enumerate(molecule["atoms"])]
+        return options
     
-    # 更新原子列表
+    # 更新原子列表显示
     def update_atoms_list(molecule):
-        atoms_data = [[atom["id"], atom["type"]] for atom in molecule["atoms"]]
+        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_id = bond["atom1"]
-            atom2_id = bond["atom2"]
-            atom1_type = next((a["type"] for a in molecule["atoms"] if a["id"] == atom1_id), "?")
-            atom2_type = next((a["type"] for a in molecule["atoms"] if a["id"] == atom2_id), "?")
+            atom1_type = molecule["atoms"][bond["atom1"]]["type"]
+            atom2_type = molecule["atoms"][bond["atom2"]]["type"]
             bonds_data.append([
-                f"{atom1_id}: {atom1_type}",
-                f"{atom2_id}: {atom2_type}",
+                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],
@@ -436,38 +437,57 @@ with gr.Blocks(title="CrystalGAT") as demo:
         inputs=molecule_state,
         outputs=atoms_list
     ).then(
-        fn=update_atom_dropdowns,
+        fn=update_atom_options,
         inputs=molecule_state,
-        outputs=[atom1_select, atom2_select]
+        outputs=atom_options
     ).then(
-        fn=lambda: "原子添加成功!",
-        outputs=status_msg
+        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)) if atom1 and atom2 else (mol, "请先选择两个原子!"),
+        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, status_msg]
+        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, bonds_list]
+        fn=lambda: [],
+        outputs=atoms_list
     ).then(
-        fn=lambda: (gr.Dropdown.update(choices=[]), gr.Dropdown.update(choices=[])),
+        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: "已清除所有原子和键",
-        outputs=status_msg
+        fn=lambda: None,
+        outputs=molecule_img
+    ).then(
+        fn=lambda: "",
+        outputs=molecule_smiles
     )
     
+    # 事件处理 - 生成分子
     generate_btn.click(
         fn=generate_smiles,
         inputs=molecule_state,
@@ -476,9 +496,6 @@ with gr.Blocks(title="CrystalGAT") as demo:
         fn=visualize_molecule,
         inputs=molecule_state,
         outputs=molecule_img
-    ).then(
-        fn=lambda: "分子生成完成!",
-        outputs=status_msg
     )
     
     # 设置交互