{ "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": [ "
" ], "text/plain": [ "" ] }, "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\"{c.get('smiles', c.get('text', 'Unknown'))}[{c.get('role', 'Unknown')}]\"\n", " for c in reaction.get(\"condition\", [])\n", " ]\n", " conditions_1 = [\n", " f\"{c.get('smiles', c.get('text', 'Unknown'))}[{c.get('role', 'Unknown')}]\"\n", " for c in reaction.get(\"condition\", [])\n", " ]\n", " products = [f\"{p.get('smiles', 'Unknown')}\" for p in reaction.get(\"products\", [])]\n", " products_1 = [f\"{p.get('smiles', 'Unknown')}\" 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\"{full_reaction}\"\n", " \n", " # 详细反应格式化输出\n", " reaction_output = f\"Reaction: {reaction_id}
\"\n", " reaction_output += f\" Reactants: {', '.join(reactants)}
\"\n", " reaction_output += f\" Conditions: {', '.join(conditions)}
\"\n", " reaction_output += f\" Products: {', '.join(products)}
\"\n", " reaction_output += f\" additional_info: {', '.join(additional_str)}
\"\n", " reaction_output += f\" Full Reaction: {full_reaction}
\"\n", " reaction_output += \"
\"\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", "

ChemEagle: A Multi-Agent System for Multimodal Chemical Information Extraction

\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 }