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