Transpolymer2 / prediction.py
transpolymer's picture
Upload prediction.py
8dad517 verified
raw
history blame
3.13 kB
import streamlit as st
import requests
from rdkit import Chem
import datetime
from db import get_database
# Function to validate SMILES string
def is_valid_smiles(smiles):
""" Validate if the input is a valid SMILES string using RDKit """
mol = Chem.MolFromSmiles(smiles)
return mol is not None
# Function to save prediction to MongoDB
def save_to_db(smiles_input, predictions):
db = get_database()
collection = db["polymers"] # your collection
doc = {
"smiles": smiles_input,
"predictions": predictions,
"timestamp": datetime.datetime.utcnow()
}
collection.insert_one(doc)
# Streamlit page for Polymer Property Prediction
def show():
st.markdown("<h1 style='text-align: center; color: #4CAF50;'>πŸ”¬ Polymer Property Prediction</h1>", unsafe_allow_html=True)
st.markdown("<hr style='border: 1px solid #ccc;'>", unsafe_allow_html=True)
# Input box with placeholder inside
input_text = st.text_input(
label="",
placeholder="Enter SMILES ",
key="smiles_input"
)
# Show Predict button always
predict_clicked = st.button("πŸš€ Predict", use_container_width=True)
# Predict on button click OR on pressing Enter with input
if (predict_clicked or input_text) and input_text.strip():
with st.spinner("Predicting..."):
if is_valid_smiles(input_text.strip()):
try:
input_data = {
"smiles": input_text.strip() # Only sending smiles
}
response = requests.post("http://127.0.0.1:8000/predict", json=input_data)
if response.status_code == 200:
result = response.json()
renamed_properties = {
"property1": "Tensile_strength (Mpa)",
"property2": "Ionization_Energy (eV)",
"property3": "Electron_Affinity (eV)",
"property4": "LogP",
"property5": "Refractive_Index",
"property6": "Molecular_Weight (g/mol)"
}
predictions = {}
for key, name in renamed_properties.items():
value = result.get(key, 'N/A')
st.markdown(f"<div style='font-size:18px; padding: 6px 0;'><strong>{name}:</strong> {value}</div>", unsafe_allow_html=True)
predictions[name] = value
# Save prediction to MongoDB
save_to_db(input_text.strip(), predictions)
st.success("Prediction saved successfully!")
else:
st.error("Prediction failed. Please try again.")
except Exception as e:
st.error(f"Error: {e}")
else:
st.error("❌ Invalid SMILES input. Please enter a correct SMILES string.")