Update app.py

app.py

@@ -212,12 +212,10 @@ def init_molecule():
     }
 
 # 添加原子
-def add_atom(molecule, atom_type, x, y):
+def add_atom(molecule, atom_type):
     molecule["atoms"].append({
         "id": len(molecule["atoms"]),
-        "type": atom_type,
-        "x": x,
-        "y": y
+        "type": atom_type
     })
     return molecule
 
@@ -303,129 +301,6 @@ def visualize_molecule(molecule_json):
         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)
-        
-        # 原子位置存储
-        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
-        )
-    
-    return interface
-
-# 查找最近的原子
-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")
@@ -451,17 +326,67 @@ with gr.Blocks(title="CrystalGAT") as demo:
     
     with gr.Tab("构建分子"):
         gr.Markdown("### 构建分子结构")
-        gr.Markdown("1. 选择原子类型和键类型")
-        gr.Markdown("2. 点击'添加原子'按钮后在画布上点击放置原子")
-        gr.Markdown("3. 点击'添加键'按钮后依次点击两个原子连接它们")
-        gr.Markdown("4. 完成后点击'生成分子'按钮")
+        gr.Markdown("1. 添加原子：选择原子类型并点击'添加原子'按钮")
+        gr.Markdown("2. 添加键：选择两个原子和键类型，然后点击'添加键'按钮")
+        gr.Markdown("3. 完成后点击'生成分子'按钮预览并预测")
+        
+        # 初始化分子状态
+        molecule_state = gr.State(init_molecule())
+        
+        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("### 添加键")
+                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=["number", "number", "str"],
+                    interactive=False
+                )
+        
+        # 操作按钮
+        with gr.Row():
+            clear_btn = gr.Button("清除所有")
+            generate_btn = gr.Button("生成分子")
         
-        # 分子构建界面
-        build_interface = build_molecule_interface()
+        # 分子预览
+        molecule_img = gr.Image(label="分子预览", interactive=False)
+        molecule_smiles = gr.Textbox(label="生成的SMILES", interactive=False)
         
         # 预测按钮
         submit_btn2 = gr.Button("使用此分子进行预测", variant="primary")
-        build_smiles = gr.Textbox(visible=False)  # 隐藏的SMILES存储
     
     # 输出区域 (两种输入方式共享)
     gr.Markdown("## 预测结果")
@@ -475,6 +400,75 @@ 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_ids = [f"{i}: {atom['type']}" for i, atom in enumerate(molecule["atoms"])]
+        return gr.Dropdown.update(choices=atom_ids), gr.Dropdown.update(choices=atom_ids)
+    
+    # 更新原子列表
+    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_dropdowns,
+        inputs=molecule_state,
+        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, bonds_list]
+    ).then(
+        fn=lambda: (gr.Dropdown.update(choices=[]), gr.Dropdown.update(choices=[])),
+        outputs=[atom1_select, atom2_select]
+    )
+    
+    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(
@@ -486,7 +480,7 @@ with gr.Blocks(title="CrystalGAT") as demo:
     # 分子构建预测路径
     submit_btn2.click(
         fn=lambda smiles: predict_all(smiles) if smiles else ("请输入有效的SMILES", None, "", None, "", None),
-        inputs=build_smiles,
+        inputs=molecule_smiles,
         outputs=[elastic_text, elastic_img, plastic_text, plastic_img, brittle_text, brittle_img]
     )