|
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 |
|
|
|
REPO_ID = "yujieq/MolScribe" |
|
FILENAME = "swin_base_char_aux_1m.pth" |
|
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) |
|
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): |
|
prediction = model.predict_image(image) |
|
smiles = prediction['smiles'] |
|
molfile = prediction['molfile'] |
|
image = generate_mol_image(molfile) |
|
return image, smiles, molfile |
|
|
|
|
|
iface = gr.Interface( |
|
predict, |
|
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. (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, |
|
) |
|
iface.launch() |
|
|