update MolScribe interface

app.py

@@ -18,6 +18,7 @@ ckpt_path = hf_hub_download(REPO_ID, FILENAME)
 device = torch.device('cpu')
 model = MolScribe(ckpt_path, device)
 
+
 def generate_mol_image(molblock):
     indigo = Indigo()
     render = IndigoRenderer(indigo)
@@ -30,10 +31,14 @@ def generate_mol_image(molblock):
     img = cv2.imdecode(np.asarray(bytearray(buf), dtype=np.uint8), 1)
     return img
 
+
 def predict(image):
-    smiles, molblock = model.predict_image(image)
-    image = generate_mol_image(molblock)
-    return image, smiles, molblock
+    prediction = model.predict_image(image)
+    smiles = prediction['smiles']
+    molfile = prediction['molfile']
+    image = generate_mol_image(molfile)
+    return image, smiles, molfile
+
 
 iface = gr.Interface(
     predict,
@@ -45,12 +50,12 @@ iface = gr.Interface(
     ],
     allow_flagging="auto",
     title="MolScribe",
-    description="Convert a molecular image into SMILES and Molfile. (It typically takes 2-3 seconds to predict an image, "
-                "but may take longer if the server is busy. "
-                "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_.",
+    description="Convert a molecular image into SMILES and Molfile. (It typically takes 2-3 seconds to predict an "
+                "image, but may take longer if the server is busy. 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,
 )

requirements.txt

@@ -1,2 +1,5 @@
 git+https://github.com/thomas0809/MolScribe.git@main#egg=MolScribe
-epam.indigo==1.8.0
+epam.indigo==1.8.0
+transformers>=4.5.1
+huggingface-hub>=0.11.0
+matplotlib>=3.5.3