| import gradio as gr | |
| import os | |
| import glob | |
| import torch | |
| from molscribe import MolScribe | |
| from huggingface_hub import hf_hub_download | |
| REPO_ID = "yujieq/MolScribe" | |
| FILENAME = "swin_base_char_aux_200k.pth" | |
| ckpt_path = hf_hub_download(REPO_ID, FILENAME) | |
| device = torch.device('cpu') | |
| model = MolScribe(ckpt_path, device) | |
| def predict(image): | |
| smiles, molblock = model.predict_image(image) | |
| return smiles, molblock | |
| iface = gr.Interface( | |
| predict, | |
| inputs=gr.Image(label="Upload molecular image"), | |
| outputs=[ | |
| gr.Textbox(label="SMILES"), | |
| gr.Textbox(label="Molfile"), | |
| ], | |
| allow_flagging="auto", | |
| title="MolScribe", | |
| description="Convert a molecular image into SMILES and Molfile. Code: https://github.com/thomas0809/MolScribe", | |
| examples=sorted(glob.glob('examples/*.png')), | |
| examples_per_page=20, | |
| ) | |
| iface.launch() |