add indigo render

app.py

@@ -2,8 +2,12 @@ import gradio as gr
 
 import os
 import glob
+import cv2
+import numpy as np
 import torch
 from molscribe import MolScribe
+from indigo import Indigo
+from indigo.renderer import IndigoRenderer
 
 from huggingface_hub import hf_hub_download
 
@@ -14,21 +18,38 @@ ckpt_path = hf_hub_download(REPO_ID, FILENAME)
 device = torch.device('cpu')
 model = MolScribe(ckpt_path, device)
 
+def generate_image(molblock):
+    indigo = Indigo()
+    render = IndigoRenderer(indigo)
+    indigo.setOption('render-output-format', 'png')
+    indigo.setOption('render-background-color', '1,1,1')
+    indigo.setOption('render-stereo-style', 'none')
+    indigo.setOption('render-label-mode', 'hetero')
+    mol = indigo.loadMolecule(molblock)
+    buf = render.renderToBuffer(mol)
+    img = cv2.imdecode(np.asarray(bytearray(buf), dtype=np.uint8), 1)
+    return img
+
 def predict(image):
     smiles, molblock = model.predict_image(image)
-    return smiles, molblock
+    image = generate_image(molblock)
+    return image, smiles, molblock
 
 iface = gr.Interface(
     predict,
-    inputs=gr.Image(label="Upload molecular image"),
+    inputs=gr.Image(label="Upload molecular image", show_label=False).style(height=256),
     outputs=[
+        gr.Image(label="Prediction").style(height=256),
         gr.Textbox(label="SMILES"),
         gr.Textbox(label="Molfile"),
     ],
     allow_flagging="auto",
     title="MolScribe",
-    description="Convert a molecular image into SMILES and Molfile. Code: https://github.com/thomas0809/MolScribe",
+    description="Convert a molecular image into SMILES and Molfile. (To view the prediction better, copy-paste the Molfile to ChemDraw.) <br> " \
+        "Paper: [_MolScribe: Robust Molecular Structure Recognition with Image-To-Graph Generation_](https://arxiv.org/abs/2205.14311) <br>" \
+        "Code: https://github.com/thomas0809/MolScribe <br>" \
+        "Authors: [Yujie Qian](mailto:[email protected]), Jiang Guo, Zhengkai Tu, Zhening Li, Connor W. Coley, Regina Barzilay. _MIT CSAIL_.",
     examples=sorted(glob.glob('examples/*.png')),
     examples_per_page=20,
 )
-iface.launch()
+iface.launch()

requirements.txt

@@ -1 +1,2 @@
 git+https://github.com/thomas0809/MolScribe.git@main#egg=MolScribe
+epam.indigo==1.8.0