veltnix/ScentNet

1import gradio as gr2import joblib3import numpy as np4import pandas as pd5import pubchempy as pcp6from rdkit import Chem7from rdkit.Chem import Draw, AllChem, Descriptors8from PIL import Image9from sklearn.preprocessing import StandardScaler10 11# --- Configuration & Constants ---12MODEL_PATH = 'saved_model/xgboost_odor_model.joblib'13ODOR_CATEGORIES = [14    'sweet', 'nutty', 'pungent', 'floral', 'fruity', 'woody', 'minty'15]16FINGERPRINT_RADIUS = 217FINGERPRINT_NBITS = 204818 19# --- Helper Functions (Copied from utils/data_utils.py for a self-contained app) ---20# We copy these functions directly into the app file to make deployment easier.21 22def smiles_to_fingerprint(smiles_string, radius=FINGERPRINT_RADIUS, n_bits=FINGERPRINT_NBITS):23    """Converts a SMILES string to a Morgan Fingerprint."""24    try:25        mol = Chem.MolFromSmiles(smiles_string)26        if mol is None: return np.zeros((n_bits,), dtype=int)27        morgan_gen = AllChem.GetMorganGenerator(radius=radius, fpSize=n_bits)28        return morgan_gen.GetFingerprintAsNumPy(mol)29    except:30        return np.zeros((n_bits,), dtype=int)31 32def calculate_physicochemical_descriptors(smiles_string):33    """Calculates a set of physicochemical descriptors for a given SMILES string."""34    try:35        mol = Chem.MolFromSmiles(smiles_string)36        if mol is None: return None37        return {38            'MolLogP': Descriptors.MolLogP(mol), 'NumHAcceptors': Descriptors.NumHAcceptors(mol),39            'NumHDonors': Descriptors.NumHDonors(mol), 'Tpsa': Descriptors.TPSA(mol),40            'MolWt': Descriptors.MolWt(mol), 'NumAromaticRings': Descriptors.NumAromaticRings(mol),41            'NumAliphaticRings': Descriptors.NumAliphaticRings(mol), 'fr_ester': Descriptors.fr_ester(mol),42            'BalabanJ': Descriptors.BalabanJ(mol),43        }44    except:45        return None46 47def create_feature_vector(smiles: str) -> np.ndarray | None:48    """Creates the complete feature vector (fingerprint + descriptors) for a single SMILES."""49    fingerprint = smiles_to_fingerprint(smiles)50    descriptors = calculate_physicochemical_descriptors(smiles)51    if not np.any(fingerprint) or descriptors is None:52        return None53    desc_df = pd.DataFrame([descriptors])54    scaler = StandardScaler()55    scaled_descriptors = scaler.fit_transform(desc_df)56    return np.hstack([fingerprint, scaled_descriptors.flatten()])57 58def get_smiles_from_name(name: str) -> str | None:59    """Looks up a chemical name on PubChem and returns its SMILES string."""60    try:61        compounds = pcp.get_compounds(name, 'name')62        return compounds[0].smiles if compounds else None63    except:64        return None65 66def draw_molecule(smiles: str) -> Image.Image | None:67    """Generates a PIL Image of a molecule's 2D structure."""68    try:69        mol = Chem.MolFromSmiles(smiles)70        return Draw.MolToImage(mol, size=(350, 200)) if mol else None71    except:72        return None73 74# --- Load the Model (globally, once) ---75print("Loading XGBoost model...")76try:77    model = joblib.load(MODEL_PATH)78    print("Model loaded successfully.")79except Exception as e:80    print(f"Error loading model: {e}")81    model = None82 83# --- The Main Prediction Function for Gradio ---84def predict_odor(input_type: str, chemical_input: str) -> (Image.Image, str):85    """86    This is the core function that Gradio will call. It takes the UI inputs,87    processes them, and returns the outputs for the UI.88    """89    if not chemical_input:90        return None, "Please enter a chemical name or SMILES string."91    if model is None:92        return None, "Model could not be loaded. Please check server logs."93 94    # Step 1: Get SMILES string95    if input_type == "Chemical Name":96        smiles = get_smiles_from_name(chemical_input)97        if smiles is None:98            return None, f"Could not find SMILES for '{chemical_input}'. Please check the name."99    else:100        smiles = chemical_input101 102    # Step 2: Generate molecule image103    structure_image = draw_molecule(smiles)104    if structure_image is None:105        return None, f"Invalid SMILES string provided: '{smiles}'"106 107    # Step 3: Create the feature vector for the model108    feature_vector = create_feature_vector(smiles)109    if feature_vector is None:110        return structure_image, "Could not generate features for the molecule."111 112    # Step 4: Make the prediction113    try:114        feature_vector_2d = feature_vector.reshape(1, -1)115        pred_proba = model.predict_proba(feature_vector_2d)[0]116        117        # Get the top two predictions118        top_2_indices = np.argsort(pred_proba)[-2:][::-1]119        120        pred_1_odor = ODOR_CATEGORIES[top_2_indices[0]]121        pred_1_confidence = pred_proba[top_2_indices[0]]122        123        pred_2_odor = ODOR_CATEGORIES[top_2_indices[1]]124        pred_2_confidence = pred_proba[top_2_indices[1]]125        126        result_text = (127            f"1. Top Prediction: {pred_1_odor.capitalize()} (Confidence: {pred_1_confidence:.2%})\n"128            f"2. Second Likelihood: {pred_2_odor.capitalize()} (Confidence: {pred_2_confidence:.2%})"129        )130        return structure_image, result_text131    except Exception as e:132        return structure_image, f"An error occurred during prediction: {e}"133 134# --- Build the Gradio Interface ---135demo = gr.Interface(136    fn=predict_odor,137    inputs=[138        gr.Radio(["SMILES String", "Chemical Name"], value="Chemical Name", label="Input Type"),139        gr.Textbox(label="Enter Input", placeholder="e.g., Vanillin or C8H8O3")140    ],141    outputs=[142        gr.Image(label="Chemical Structure", type="pil"),143        gr.Textbox(label="Result:")144    ],145    title="ScentNet - Chemical Odor Predictor",146    description="Enter a chemical's name or its SMILES string to predict its primary odor category. This demo uses an XGBoost model trained on molecular fingerprints and physicochemical descriptors.",147    examples=[148        ["Chemical Name", "Vanillin"],149        ["Chemical Name", "Benzaldehyde"],150        ["SMILES String", "CCO"],151        ["SMILES String", "C1=CC=C(C=C1)CC(=O)C"]152    ],153    allow_flagging="never"154)155 156# --- Launch the App ---157if __name__ == "__main__":158    demo.launch()159