Update app.py

app.py

@@ -8,7 +8,9 @@ from rxn.reaction import Reaction
 from rdkit import Chem
 from rdkit.Chem import rdChemReactions
 from rdkit.Chem import Draw
-
+from rdkit.Chem import AllChem
+from rdkit.Chem.Draw import rdMolDraw2D
+import cairosvg
 
 PROMPT_DIR = "prompts/"
 ckpt_path = "./rxn/model/model.ckpt"
@@ -19,7 +21,7 @@ PROMPT_NAMES = {
     "2_RxnOCR.txt": "Reaction Image Parsing Workflow",
 }
 example_diagram = "examples/exp.png"
-rdkit_image = "examples/image.webp"
+rdkit_image = "examples/rdkit.png"
 
 def list_prompt_files_with_names():
     """
@@ -106,10 +108,10 @@ def process_chem_image(image, selected_task):
     return "\n\n".join(detailed_reactions), smiles_output, combined_image_path, example_diagram, json_file_path
 
 
-
+# 获取 prompts 和友好名字
 prompts_with_names = list_prompt_files_with_names()
 
-
+# 示例数据：图像路径 + 任务选项
 examples = [
     
     ["examples/reaction1.png", "Reaction Image Parsing Workflow"],
@@ -167,7 +169,7 @@ with gr.Blocks() as demo:
             @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= "SMILES of Reaction i"), gr.Image(value=rdkit_image, label= "RDKit Image of Reaction i")
+                    return gr.Textbox(label= "SMILES of Reaction i"), gr.Image(value=rdkit_image, label= "RDKit Image of Reaction i",height=100)
                 else:
                     # 假设输入是逗号分隔的 SMILES 字符串
                     smiles_list = inputs.split(",")
@@ -175,14 +177,65 @@ with gr.Blocks() as demo:
                     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"SMILES of Reaction {i + 1} ", show_copy_button=True, interactive=False))
-                        components.append(gr.Image(value=img,label= f"RDKit Image of Reaction {i + 1} ")) 
+                        rxn = rdChemReactions.ReactionFromSmarts(smiles, useSmiles=True)
+                        
+                        if rxn:
+
+                            new_rxn = AllChem.ChemicalReaction()	
+                            for mol in rxn.GetReactants():
+                                mol = Chem.MolFromMolBlock(Chem.MolToMolBlock(mol))
+                                new_rxn.AddReactantTemplate(mol)
+                            for mol in rxn.GetProducts():
+                                mol = Chem.MolFromMolBlock(Chem.MolToMolBlock(mol))
+                                new_rxn.AddProductTemplate(mol)
+
+                            rxn = new_rxn
+
+                            def atom_mapping_remover(rxn):
+                                for reactant in rxn.GetReactants():
+                                    for atom in reactant.GetAtoms():
+                                        atom.SetAtomMapNum(0)
+                                for product in rxn.GetProducts():
+                                    for atom in product.GetAtoms():
+                                        atom.SetAtomMapNum(0)
+                                return rxn
+                            
+                            atom_mapping_remover(rxn)
+
+                            reactant1 = rxn.GetReactantTemplate(0)
+                            print(reactant1.GetNumBonds)
+                            reactant2 = rxn.GetReactantTemplate(1) if rxn.GetNumReactantTemplates() > 1 else None
+
+                            if reactant1.GetNumBonds() > 0:
+                                bond_length_reference = Draw.MeanBondLength(reactant1)
+                            elif reactant2 and reactant2.GetNumBonds() > 0:
+                                bond_length_reference = Draw.MeanBondLength(reactant2)
+                            else:
+                                bond_length_reference = 1.0 
+
+
+                            drawer = rdMolDraw2D.MolDraw2DSVG(-1, -1)
+                            dopts = drawer.drawOptions()
+                            dopts.padding = 0.1 
+                            dopts.includeRadicals = True
+                            Draw.SetACS1996Mode(dopts, bond_length_reference*0.55)
+                            dopts.bondLineWidth = 1.5
+                            drawer.DrawReaction(rxn)
+                            drawer.FinishDrawing()
+                            svg_content = drawer.GetDrawingText()
+                            svg_file = f"reaction{i+1}.svg"
+                            with open(svg_file, "w") as f:
+                                f.write(svg_content)
+                            png_file = f"reaction_{i+1}.png"
+                            cairosvg.svg2png(url=svg_file, write_to=png_file)
+
+
+                        
+                        components.append(gr.Textbox(value=smiles,label= f"Reaction {i + 1} SMILES", show_copy_button=True, interactive=False))
+                        components.append(gr.Image(value=png_file,label= f"Reaction {i + 1} RDKit Image")) 
                     return components  # 返回包含所有 SMILES Textbox 组件的列表
                 
-            download_json = gr.File(label="Download JSON File")
+            download_json = gr.File(label="Download JSON File",)