Update app.py

app.py

@@ -5,6 +5,10 @@ from rxnim import RXNIM
 from getReaction import generate_combined_image
 import torch
 from rxn.reaction import Reaction
+from rdkit import Chem
+from rdkit.Chem import rdChemReactions
+from rdkit.Chem import Draw
+
 
 PROMPT_DIR = "prompts/"
 ckpt_path = "./rxn/model/model.ckpt"
@@ -15,6 +19,7 @@ PROMPT_NAMES = {
     "2_RxnOCR.txt": "Reaction Image Parsing Workflow",
 }
 example_diagram = "examples/exp.png"
+rdkit_image = "examples/image.webp"
 
 def list_prompt_files_with_names():
     """
@@ -36,6 +41,7 @@ def parse_reactions(output_json):
     reactions_data = json.loads(output_json)  # 转换 JSON 字符串为字典
     reactions_list = reactions_data.get("reactions", [])
     detailed_output = []
+    smiles_output = [] 
 
     for reaction in reactions_list:
         reaction_id = reaction.get("reaction_id", "Unknown ID")
@@ -50,6 +56,7 @@ def parse_reactions(output_json):
         ]
         products = [f"<span style='color:orange'>{p.get('smiles', 'Unknown')}</span>" for p in reaction.get("products", [])]
         products_1 = [f"<span style='color:black'>{p.get('smiles', 'Unknown')}</span>" for p in reaction.get("products", [])]
+        products_2 = [r.get("smiles", "Unknown") for r in reaction.get("products", [])]
 
         # 构造反应的完整字符串，定制字体颜色
         full_reaction = f"{'.'.join(reactants)}>>{'.'.join(products_1)} | {', '.join(conditions_1)}"
@@ -64,7 +71,12 @@ def parse_reactions(output_json):
         reaction_output += "<br>"
         detailed_output.append(reaction_output)
 
-    return detailed_output
+        reaction_smiles = f"{'.'.join(reactants)}>>{'.'.join(products_2)}"
+        smiles_output.append(reaction_smiles)
+
+
+
+    return detailed_output, smiles_output
 
 def process_chem_image(image, selected_task):
     chem_mllm = RXNIM()
@@ -78,11 +90,12 @@ def process_chem_image(image, selected_task):
     rxnim_result = chem_mllm.process(image_path, prompt_path)
 
     # 将 JSON 结果解析为结构化输出
-    detailed_reactions = parse_reactions(rxnim_result)
+    detailed_reactions, smiles_output = parse_reactions(rxnim_result)
 
     # 调用 RxnScribe 模型处理并生成整合图像
     predictions = model.predict_image_file(image_path, molscribe=True, ocr=True)
     combined_image_path = generate_combined_image(predictions, image_path)
+    #combined_image_path = model.draw_predictions(predictions, image_path)
 
     json_file_path = "output.json"
     with open(json_file_path, "w") as json_file:
@@ -90,7 +103,7 @@ def process_chem_image(image, selected_task):
 
 
     # 返回详细反应和整合图像
-    return "\n\n".join(detailed_reactions), combined_image_path, example_diagram, json_file_path
+    return "\n\n".join(detailed_reactions), smiles_output, combined_image_path, example_diagram, json_file_path
 
 
 # 获取 prompts 和友好名字
@@ -106,26 +119,111 @@ examples = [
 ]
 
 # 定义 Gradio 界面
-demo = gr.Interface(
-    fn=process_chem_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Reaction Image"),
-        gr.Radio(
-            choices=list(prompts_with_names.keys()),  # 显示任务名字
-            label="Select a predefined task",
-        ),
-    ],
-    outputs=[
-        gr.HTML(label="Reaction outputs"),
-        gr.Image(label="Visualization"), # 显示整合图像
-        gr.Image(value=example_diagram, label="Schematic Diagram"),
-        gr.File(label="Download JSON File"),
-
-    ],
-    title="Towards Large-scale Chemical Reaction Image Parsing via a Multimodal Large Language Model",
-    description="Upload a reaction image and select a predefined task prompt.",
-    examples=examples,  # 使用嵌套列表作为示例
-    examples_per_page=20,
-)
-
-demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown(
+"""
+
+    <center> <h1>Towards Large-scale Chemical Reaction Image Parsing via a Multimodal Large Language Model<h1></center>
+
+    Upload a reaction image and select a predefined task prompt.
+    """)
+
+    
+
+    # 上半部分，输入区域
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1):  # 左侧列
+            image_input = gr.Image(type="pil", label="Upload Reaction Image")
+            task_radio = gr.Radio(
+                choices=list(prompts_with_names.keys()),
+                label="Select a predefined task",
+            )
+            with gr.Row():  # Clear 和 Submit 按钮放在同一行
+                clear_button = gr.Button("Clear")
+                process_button = gr.Button("Run", elem_id="submit-btn")
+            
+            gr.Markdown("### Reaction Imge Parsing Output")
+            reaction_output = gr.HTML(label="Reaction outputs")
+            
+        
+        with gr.Column(scale=1):
+                
+                gr.Markdown("### Reaction Extraction Output")
+                visualization_output = gr.Image(label="Visualization Output") 
+                schematic_diagram = gr.Image(value=example_diagram, label="Schematic Diagram")      
+                
+
+        with gr.Column(scale=1):  
+            gr.Markdown("### Machine-readable Data Output")
+            smiles_output = gr.Textbox(
+                label="Reaction SMILES",
+                show_copy_button=True,
+                interactive=False,
+                visible=False,
+            )
+            
+
+    # 下半部分，图像和 JSON 输出
+            @gr.render(inputs = smiles_output)  # 使用gr.render修饰器绑定输入和渲染逻辑
+            def show_split(inputs):  # 定义处理和展示分割文本的函数
+                if not inputs or isinstance(inputs, str) and inputs.strip() == "":  # 检查输入文本是否为空
+                    return gr.Textbox(label= f"Reaction SMILES"), gr.Image(value=rdkit_image, label= "RDKit Image generated from Reaction SMILES")
+                else:
+                    # 假设输入是逗号分隔的 SMILES 字符串
+                    smiles_list = inputs.split(",")
+                    smiles_list = [item.strip("[]' ") for item in smiles_list]
+                    components = []  # 初始化一个组件列表，用于存放每个 SMILES 对应的 Textbox 组件
+                    for i, smiles in enumerate(smiles_list): 
+                        smiles.replace('"', '').replace("'", "").replace("[", "").replace("]", "")
+                        reaction = rdChemReactions.ReactionFromSmarts(smiles)
+                        if reaction:
+                            img = Draw.ReactionToImage(reaction)
+                        components.append(gr.Textbox(value=smiles,label= f"Reaction {i + 1} SMILES", show_copy_button=True, interactive=False))
+                        components.append(gr.Image(value=img,label= f"Reaction {i + 1} RDKit Image")) 
+                    return components  # 返回包含所有 SMILES Textbox 组件的列表
+                
+            download_json = gr.File(label="Download JSON File",)
+        
+        
+            
+
+    # 示例部分
+    gr.Examples(
+        examples=examples,
+        inputs=[image_input, task_radio],
+        outputs=[reaction_output, smiles_output, visualization_output],
+    )
+
+    # 绑定功能
+    clear_button.click(
+        lambda: (None, None, None, None, None),
+        inputs=[],
+        outputs=[
+            image_input,
+            task_radio,
+            reaction_output,
+            smiles_output,
+            visualization_output,
+        ],
+    )
+
+    process_button.click(
+        process_chem_image,
+        inputs=[image_input, task_radio],
+        outputs=[
+            reaction_output,
+            smiles_output,
+            visualization_output,
+            schematic_diagram,
+            download_json,
+        ],
+    )
+
+demo.css = """
+#submit-btn {
+    background-color: #FF914D;
+    color: white;
+    font-weight: bold;
+}
+"""
+demo.launch()