Spaces:
Running
Running
File size: 2,226 Bytes
866233d e09f96a bec3c91 e09f96a bec3c91 866233d e09f96a f07373b e09f96a 136efdc f07373b bec3c91 136efdc e09f96a 136efdc e09f96a bec3c91 e09f96a bec3c91 d30c1c3 e09f96a 136efdc e09f96a bec3c91 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
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").style(show_copy_button=True),
gr.Textbox(label="Molfile").style(show_copy_button=True),
],
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()
|