File size: 14,066 Bytes
1f516b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13d3631",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7866\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7866/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import gradio as gr\n",
    "import json\n",
    "from main import ChemEagle  # 支持 API key 通过环境变量\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import rdChemReactions\n",
    "from rdkit.Chem import Draw\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit.Chem.Draw import rdMolDraw2D\n",
    "import cairosvg\n",
    "import re\n",
    "import torch\n",
    "\n",
    "example_diagram = \"examples/exp.png\"\n",
    "rdkit_image = \"examples/rdkit.png\"\n",
    "# 解析 ChemEagle 返回的结构化数据\n",
    "def parse_reactions(output_json):\n",
    "    \"\"\"\n",
    "    解析 JSON 格式的反应数据并格式化输出,包含颜色定制。\n",
    "    \"\"\"\n",
    "    if isinstance(output_json, str):\n",
    "        reactions_data = json.loads(output_json)\n",
    "    elif isinstance(output_json, dict):\n",
    "        reactions_data = output_json  # 转换 JSON 字符串为字典\n",
    "    reactions_list = reactions_data.get(\"reactions\", [])\n",
    "    detailed_output = []\n",
    "    smiles_output = [] \n",
    "\n",
    "    for reaction in reactions_list:\n",
    "        reaction_id = reaction.get(\"reaction_id\", \"Unknown ID\")\n",
    "        reactants = [r.get(\"smiles\", \"Unknown\") for r in reaction.get(\"reactants\", [])]\n",
    "        conditions = [\n",
    "            f\"<span style='color:red'>{c.get('smiles', c.get('text', 'Unknown'))}[{c.get('role', 'Unknown')}]</span>\"\n",
    "            for c in reaction.get(\"condition\", [])\n",
    "        ]\n",
    "        conditions_1 = [\n",
    "            f\"<span style='color:black'>{c.get('smiles', c.get('text', 'Unknown'))}[{c.get('role', 'Unknown')}]</span>\"\n",
    "            for c in reaction.get(\"condition\", [])\n",
    "        ]\n",
    "        products = [f\"<span style='color:orange'>{p.get('smiles', 'Unknown')}</span>\" for p in reaction.get(\"products\", [])]\n",
    "        products_1 = [f\"<span style='color:black'>{p.get('smiles', 'Unknown')}</span>\" for p in reaction.get(\"products\", [])]\n",
    "        products_2 = [r.get(\"smiles\", \"Unknown\") for r in reaction.get(\"products\", [])]\n",
    "        \n",
    "        additional = reaction.get(\"additional_info\", [])\n",
    "        additional_str = [str(x) for x in additional if x is not None]\n",
    "\n",
    "        tail = conditions_1 + additional_str\n",
    "        tail_str = \", \".join(tail)\n",
    "\n",
    "        # 构造反应的完整字符串,定制字体颜色\n",
    "        full_reaction = f\"{'.'.join(reactants)}>>{'.'.join(products_1)} | {tail_str}\"\n",
    "        full_reaction = f\"<span style='color:black'>{full_reaction}</span>\"\n",
    "        \n",
    "        # 详细反应格式化输出\n",
    "        reaction_output = f\"<b>Reaction: </b> {reaction_id}<br>\"\n",
    "        reaction_output += f\"  Reactants: <span style='color:blue'>{', '.join(reactants)}</span><br>\"\n",
    "        reaction_output += f\"  Conditions: {', '.join(conditions)}<br>\"\n",
    "        reaction_output += f\"  Products: {', '.join(products)}<br>\"\n",
    "        reaction_output += f\"  additional_info: {', '.join(additional_str)}<br>\"\n",
    "        reaction_output += f\"  <b>Full Reaction:</b> {full_reaction}<br>\"\n",
    "        reaction_output += \"<br>\"\n",
    "        detailed_output.append(reaction_output)\n",
    "\n",
    "        reaction_smiles = f\"{'.'.join(reactants)}>>{'.'.join(products_2)}\"\n",
    "        smiles_output.append(reaction_smiles)\n",
    "    return detailed_output, smiles_output\n",
    "\n",
    "\n",
    "# 核心处理函数,仅使用 API Key 和图像\n",
    "def process_chem_image(api_key, image):\n",
    "    # 设置 API Key 环境变量,供 ChemEagle 使用\n",
    "    os.environ[\"CHEMEAGLE_API_KEY\"] = api_key\n",
    "\n",
    "    # 保存上传图片\n",
    "    image_path = \"temp_image.png\"\n",
    "    image.save(image_path)\n",
    "\n",
    "    # 调用 ChemEagle(实现内部读取 os.getenv)\n",
    "    chemeagle_result = ChemEagle(image_path)\n",
    "\n",
    "    # 解析输出\n",
    "    detailed, smiles = parse_reactions(chemeagle_result)\n",
    "\n",
    "    # 写出 JSON\n",
    "    json_path = \"output.json\"\n",
    "    with open(json_path, 'w') as jf:\n",
    "        json.dump(chemeagle_result, jf, indent=2)\n",
    "\n",
    "    # 返回 HTML、SMILES 合并文本、示意图、JSON 下载\n",
    "    return \"\\n\\n\".join(detailed), smiles, example_diagram, json_path\n",
    "\n",
    "# 构建 Gradio 界面\n",
    "with gr.Blocks() as demo:\n",
    "    gr.Markdown(\n",
    "        \"\"\"\n",
    "        <center><h1>ChemEagle: A Multi-Agent System for Multimodal Chemical Information Extraction</h1></center>\n",
    "        Upload a multimodal reaction image and type your OpenAI API key to extract multimodal chemical information.\n",
    "        \"\"\"\n",
    "    )\n",
    "\n",
    "    with gr.Row():\n",
    "        # ———— 左侧:上传 + API Key + 按钮 ————\n",
    "        with gr.Column(scale=1):\n",
    "            image_input   = gr.Image(type=\"pil\", label=\"Upload a multimodal reaction image\")\n",
    "            api_key_input = gr.Textbox(\n",
    "                label=\"Your API-Key\",\n",
    "                placeholder=\"Type your OpenAI_API_KEY\",\n",
    "                type=\"password\"\n",
    "            )\n",
    "            with gr.Row():\n",
    "                clear_btn = gr.Button(\"Clear\")\n",
    "                run_btn   = gr.Button(\"Run\", elem_id=\"submit-btn\")\n",
    "\n",
    "        # ———— 中间:解析结果 + 示意图 ————\n",
    "        with gr.Column(scale=1):\n",
    "            gr.Markdown(\"### Parsed Reactions\")\n",
    "            reaction_output   = gr.HTML(label=\"Detailed Reaction Output\")\n",
    "            gr.Markdown(\"### Schematic Diagram\")\n",
    "            schematic_diagram = gr.Image(value=example_diagram, label=\"示意图\")\n",
    "\n",
    "        # ———— 右侧:SMILES 拆分 & RDKit 渲染 + JSON 下载 ————\n",
    "        with gr.Column(scale=1):\n",
    "            gr.Markdown(\"### Machine-readable Output\")\n",
    "            smiles_output = gr.Textbox(\n",
    "                label=\"Reaction SMILES\",\n",
    "                show_copy_button=True,\n",
    "                interactive=False,\n",
    "                visible=False\n",
    "            )\n",
    "\n",
    "            @gr.render(inputs = smiles_output)  # 使用gr.render修饰器绑定输入和渲染逻辑\n",
    "            def show_split(inputs):  # 定义处理和展示分割文本的函数\n",
    "                if not inputs or isinstance(inputs, str) and inputs.strip() == \"\":  # 检查输入文本是否为空\n",
    "                    return gr.Textbox(label= \"SMILES of Reaction i\"), gr.Image(value=rdkit_image, label= \"RDKit Image of Reaction i\",height=100)\n",
    "                else:\n",
    "                    # 假设输入是逗号分隔的 SMILES 字符串\n",
    "                    smiles_list = inputs.split(\",\")\n",
    "                    smiles_list = [re.sub(r\"^\\s*\\[?'?|'\\]?\\s*$\", \"\", item) for item in smiles_list]\n",
    "                    components = []  # 初始化一个组件列表,用于存放每个 SMILES 对应的 Textbox 组件\n",
    "                    for i, smiles in enumerate(smiles_list): \n",
    "                        smiles.replace('\"', '').replace(\"'\", \"\").replace(\"[\", \"\").replace(\"]\", \"\")\n",
    "                        rxn = rdChemReactions.ReactionFromSmarts(smiles, useSmiles=True)\n",
    "                        \n",
    "                        if rxn:\n",
    "\n",
    "                            new_rxn = AllChem.ChemicalReaction()\t\n",
    "                            for mol in rxn.GetReactants():\n",
    "                                mol = Chem.MolFromMolBlock(Chem.MolToMolBlock(mol))\n",
    "                                new_rxn.AddReactantTemplate(mol)\n",
    "                            for mol in rxn.GetProducts():\n",
    "                                mol = Chem.MolFromMolBlock(Chem.MolToMolBlock(mol))\n",
    "                                new_rxn.AddProductTemplate(mol)\n",
    "\n",
    "                            rxn = new_rxn\n",
    "\n",
    "                            def atom_mapping_remover(rxn):\n",
    "                                for reactant in rxn.GetReactants():\n",
    "                                    for atom in reactant.GetAtoms():\n",
    "                                        atom.SetAtomMapNum(0)\n",
    "                                for product in rxn.GetProducts():\n",
    "                                    for atom in product.GetAtoms():\n",
    "                                        atom.SetAtomMapNum(0)\n",
    "                                return rxn\n",
    "                            \n",
    "                            atom_mapping_remover(rxn)\n",
    "\n",
    "                            reactant1 = rxn.GetReactantTemplate(0)\n",
    "                            print(reactant1.GetNumBonds)\n",
    "                            reactant2 = rxn.GetReactantTemplate(1) if rxn.GetNumReactantTemplates() > 1 else None\n",
    "\n",
    "                            if reactant1.GetNumBonds() > 0:\n",
    "                                bond_length_reference = Draw.MeanBondLength(reactant1)\n",
    "                            elif reactant2 and reactant2.GetNumBonds() > 0:\n",
    "                                bond_length_reference = Draw.MeanBondLength(reactant2)\n",
    "                            else:\n",
    "                                bond_length_reference = 1.0 \n",
    "\n",
    "\n",
    "                            drawer = rdMolDraw2D.MolDraw2DSVG(-1, -1)\n",
    "                            dopts = drawer.drawOptions()\n",
    "                            dopts.padding = 0.1 \n",
    "                            dopts.includeRadicals = True\n",
    "                            Draw.SetACS1996Mode(dopts, bond_length_reference*0.55)\n",
    "                            dopts.bondLineWidth = 1.5\n",
    "                            drawer.DrawReaction(rxn)\n",
    "                            drawer.FinishDrawing()\n",
    "                            svg_content = drawer.GetDrawingText()\n",
    "                            svg_file = f\"reaction{i+1}.svg\"\n",
    "                            with open(svg_file, \"w\") as f:\n",
    "                                f.write(svg_content)\n",
    "                            png_file = f\"reaction_{i+1}.png\"\n",
    "                            cairosvg.svg2png(url=svg_file, write_to=png_file)\n",
    "\n",
    "\n",
    "                        \n",
    "                        components.append(gr.Textbox(value=smiles,label= f\"SMILES of Reaction {i}\", show_copy_button=True, interactive=False))\n",
    "                        components.append(gr.Image(value=png_file,label= f\"RDKit Image of Reaction {i}\")) \n",
    "                    return components  # 返回包含所有 SMILES Textbox 组件的列表\n",
    "\n",
    "            download_json = gr.File(label=\"Download JSON File\")\n",
    "\n",
    "\n",
    "    gr.Examples(\n",
    "        examples=[\n",
    "            [\"examples/reaction1.jpg\", \"\"],\n",
    "            [\"examples/reaction2.png\", \"\"],\n",
    "            [\"examples/reaction3.png\", \"\"],\n",
    "            [\"examples/reaction4.png\", \"\"],\n",
    "            \n",
    "            \n",
    "        ],\n",
    "        inputs=[image_input, api_key_input],\n",
    "        outputs=[reaction_output, smiles_output, schematic_diagram, download_json],\n",
    "        cache_examples=False,\n",
    "        examples_per_page=4,\n",
    "    )\n",
    "\n",
    "    # ———— 清空与运行 绑定 ————\n",
    "    clear_btn.click(\n",
    "        lambda: (None, None, None, None, None),\n",
    "        inputs=[],\n",
    "        outputs=[image_input, api_key_input, reaction_output, smiles_output, download_json]\n",
    "    )\n",
    "    run_btn.click(\n",
    "        process_chem_image,\n",
    "        inputs=[api_key_input, image_input],\n",
    "        outputs=[reaction_output, smiles_output, schematic_diagram, download_json]\n",
    "    )\n",
    "\n",
    "    # 自定义按钮样式\n",
    "    demo.css = \"\"\"\n",
    "    #submit-btn {\n",
    "        background-color: #FF914D;\n",
    "        color: white;\n",
    "        font-weight: bold;\n",
    "    }\n",
    "    \"\"\"\n",
    "\n",
    "    demo.launch()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openchemie",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}