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()