| 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 = "Erawan/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="Internal Erawan team", |
| examples=sorted(glob.glob('examples/*.png')), |
| examples_per_page=20, |
| ) |
| iface.launch() |
|
|
