app.py

import streamlit as st
from PIL import Image
import numpy as np
import pubchempy as pcp
from rdkit import Chem
from rdkit.Chem import Draw
import time

st.title("3D2SMILES")
st.markdown("""
This app generates SMILES strings from images of molecules ball-and-stick models. 

[Version 1 Paper](https://chemrxiv.org/engage/chemrxiv/article-details/673a9d62f9980725cf89abe1) | 
[Version 2 Paper]() |
[Synthetic Dataset](https://huggingface.co/datasets/weathon/3d2smiles_synthetic) | 
[Real Dataset](https://huggingface.co/datasets/weathon/3d2smiles_real) |
[Author Github](https://github.com/weathon) |
[Feedback](mailto:[email protected]) | 
[Deploy](https://huggingface.co/spaces/weathon/3d2smiles?docker=true)
""")
col1, col2 = st.columns(2)
gen_strategy = col1.selectbox("Select a generative strategy", ("Beam Search", "Sampling", "Greedy Search"))
model = col2.selectbox("Select a model", ("V1", "V2"))
uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
if uploaded_file:
    start_time = time.time()
    image = Image.open(uploaded_file)

    options = ["CC(=O)OC1=CC=CC=C1C(=C)C(=O)O", "CC(=O)", "CC(=O)O", "CC(=O)C", "CC(=O)C1=CC=CC=C1"]
    grid = [st.columns(2) for _ in range(len(options) // 3 + 1)]
    cols = [col for row in grid for col in row]

    for i, (smiles, col) in enumerate(zip(options, cols)):
        cid = pcp.get_compounds(smiles, 'smiles')
        name = cid[0].synonyms[0]
        col.markdown("## Option {}: {}".format(i + 1, name))
        m = Chem.MolFromSmiles(smiles)
        img = Draw.MolToImage(m)
        col.image(img, use_container_width=False)
        pubchem_url = "https://pubchem.ncbi.nlm.nih.gov/compound/{}".format(cid[0].cid)
        col.markdown("[PubChem Link]({})".format(pubchem_url))
    st.markdown("---")
    st.markdown("Taken {} seconds".format(round(time.time() - start_time, 2)))