| import gradio as gr | |
| import os | |
| import glob | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from rxnscribe import RxnScribe | |
| from huggingface_hub import hf_hub_download | |
| REPO_ID = "yujieq/RxnScribe" | |
| FILENAME = "pix2seq_reaction_full.ckpt" | |
| ckpt_path = hf_hub_download(REPO_ID, FILENAME) | |
| device = torch.device('cpu') | |
| model = RxnScribe(ckpt_path, device) | |
| def get_markdown(reaction): | |
| output = [] | |
| for x in ['reactants', 'conditions', 'products']: | |
| s = '' | |
| for ent in reaction[x]: | |
| if 'smiles' in ent: | |
| s += ent['smiles'] + '<br>' | |
| elif 'text' in ent: | |
| s += ' '.join(ent['text']) + '<br>' | |
| else: | |
| s += ent['category'] | |
| output.append(s) | |
| return output | |
| def predict(image, molscribe, ocr): | |
| predictions = model.predict_image(image, molscribe=molscribe, ocr=ocr) | |
| pred_images = model.draw_predictions(predictions, image=image) | |
| markdown = [[i] + get_markdown(reaction) for i, reaction in enumerate(predictions)] | |
| return pred_images, markdown | |
| with gr.Blocks() as demo: | |
| with gr.Column(): | |
| with gr.Row(): | |
| image = gr.Image(label="Upload reaction diagram", show_label=False, type='pil').style(height=256) | |
| with gr.Row(): | |
| molscribe = gr.Checkbox(label="Run MolScribe to recognize molecule structures") | |
| ocr = gr.Checkbox(label="Run OCR to recognize text") | |
| btn = gr.Button("Submit").style(full_width=False) | |
| with gr.Row(): | |
| gallery = gr.Gallery( | |
| label="Predicted reactions", show_label=False, elem_id="gallery" | |
| ).style(height="auto") | |
| markdown = gr.Dataframe( | |
| headers=['#', 'reactant', 'condition', 'product'], | |
| datatype=['number'] + ['markdown'] * 3, | |
| wrap=False | |
| ) | |
| btn.click(predict, inputs=[image, molscribe, ocr], outputs=[gallery, markdown]) | |
| demo.launch() | |