Update app.py

app.py

@@ -1,8 +1,8 @@
-
 import os
 import zipfile
 import logging
 import torch
+import torch.nn.functional as F
 import numpy as np
 from io import BytesIO
 from PIL import Image
@@ -10,7 +10,7 @@ from PIL import Image
 import gradio as gr
 from torch_geometric.data import Data as PyGData
 import matplotlib
-matplotlib.use('Agg')  # fix backend threading issue
+matplotlib.use('Agg')
 
 from rdkit import Chem
 from rdkit.Chem import Draw, AllChem, MolFromSmiles
@@ -62,10 +62,10 @@ if torch.cuda.is_available():
 # Model Loading
 # ----------------------------
 def load_models():
-    import numpy.core.multiarray
     from torch.serialization import add_safe_globals
+    import numpy.core.multiarray
 
-    # 显式允许加载某些 numpy 序列化类型（仅当你信任模型来源）
+    # allow safe numpy objects if needed
     add_safe_globals([numpy.core.multiarray.scalar])
 
     specs = {
@@ -75,11 +75,10 @@ def load_models():
     }
 
     models = {}
-
     for name, (path, out_dim) in specs.items():
         if not os.path.exists(path):
             if os.path.exists("models.zip"):
-                logging.info("Unzipping models.zip...")
+                logger.info("Extracting models.zip...")
                 with zipfile.ZipFile("models.zip", 'r') as z:
                     z.extractall(".")
             else:
@@ -87,29 +86,30 @@ def load_models():
 
         model = EnhancedGAT(input_dim=12, hidden_dim=512, output_dim=out_dim, num_heads=8)
 
-        # 🔧 明确禁用 weights_only 以兼容新 PyTorch 版本
         try:
             state = torch.load(path, map_location=device, weights_only=False)
         except TypeError:
-            # 兼容旧版本 PyTorch（无 weights_only 参数）
             state = torch.load(path, map_location=device)
 
         state_dict = state.get("model_state_dict", state)
         model.load_state_dict(state_dict)
         model.eval().to(device)
         models[name] = model
-        logging.info(f"{name} model loaded successfully.")
+        logger.info(f"{name} model loaded successfully.")
 
     return models
 
-
 models = load_models()
 
 # ----------------------------
 # Prediction Function
 # ----------------------------
 def predict_all(smiles: str):
-    """Run Elastic, Plastic, Brittle predictions; return text & PIL images."""
+    """
+    Run predictions for Elastic, Plastic, Brittle.
+    Use threshold 0.5 for Elastic/Brittle, 0.3 for Plastic.
+    Return (text, PIL image) for each.
+    """
     atom_feats, (rows, cols, edge_attr), _ = smiles_to_graph(smiles)
     x = torch.tensor(atom_feats, dtype=torch.float)
     edge_index = torch.tensor(np.vstack((rows, cols)), dtype=torch.long)
@@ -117,17 +117,26 @@ def predict_all(smiles: str):
     data = PyGData(x=x, edge_index=edge_index, edge_attr=edge_attr,
                    smiles=[smiles], batch=torch.zeros(x.size(0), dtype=torch.long))
 
-    results = []
+    outputs = []
+    thresholds = {"Elastic": 0.5, "Plastic": 0.3, "Brittle": 0.5}
+
     for name in ["Elastic", "Plastic", "Brittle"]:
-        buf, pred = visualize_single_molecule(models[name], data, device, name)
-        if buf:
-            buf.seek(0)
-            img = Image.open(buf)
-            text = f"{name}: {int(pred)}"
-            results.append((text, img))
-        else:
-            results.append((f"{name} prediction failed", None))
-    return (*results[0], *results[1], *results[2])
+        model = models[name]
+        with torch.no_grad():
+            logits = model(data)
+            # assume binary classification: two outputs
+            if logits.dim() == 1 or logits.size(1) == 1:
+                prob = torch.sigmoid(logits).item()
+            else:
+                prob = F.softmax(logits, dim=1)[0, 1].item()
+        label = int(prob >= thresholds[name])
+        # get visualization buffer
+        buf, _ = visualize_single_molecule(model, data, device, name)
+        img = Image.open(buf) if buf else None
+        outputs.append((f"{name}: {label} (p={prob:.2f})", img))
+
+    # flatten to 6 outputs
+    return (*outputs[0], *outputs[1], *outputs[2])
 
 # ----------------------------
 # Molecule Builder Utilities
@@ -201,31 +210,31 @@ def update_bonds_list(mol):
 # ----------------------------
 with gr.Blocks(title="CrystalGAT", css="""
     .gradio-container {max-width:800px; margin:auto}
-    .block {padding:1em}
     .gr-button {margin:0.2em}
 """) as demo:
+
     gr.Markdown("## CrystalGAT  \nEnter a SMILES string or build a molecule to predict Elastic, Plastic, and Brittle classes with attention visualization.")
 
     with gr.Tab("SMILES Input"):
         smi_in = gr.Textbox(label="SMILES", placeholder="e.g. CCO")
         predict1 = gr.Button("Predict")
-    
-    with gr.Tab("Manual Molecules Construction"):
+
+    with gr.Tab("Manual Molecule Construction"):
         state = gr.State(init_molecule())
         status = gr.Textbox(label="Status", interactive=False, value="Start by adding atoms")
-        
+
         with gr.Row():
-            with gr.Column(scale=1):
+            with gr.Column():
                 atom_type = gr.Dropdown(label="Atom Type", choices=ATOM_TYPES, value="C")
                 add_a = gr.Button("Add Atom")
                 atom_tbl = gr.Dataframe(headers=["ID","Type"], datatype=["number","str"], interactive=False)
-            with gr.Column(scale=1):
+            with gr.Column():
                 a1 = gr.Dropdown(label="Atom 1", choices=[], value=None)
                 a2 = gr.Dropdown(label="Atom 2", choices=[], value=None)
                 bond_type = gr.Dropdown(label="Bond Type", choices=BOND_TYPES, value="Single")
                 add_b = gr.Button("Add Bond")
                 bond_tbl = gr.Dataframe(headers=["Atom1","Atom2","Type"], datatype=["str","str","str"], interactive=False)
-        
+
         with gr.Row():
             clear = gr.Button("Clear All")
             make = gr.Button("Generate SMILES")
@@ -234,7 +243,7 @@ with gr.Blocks(title="CrystalGAT", css="""
 
         predict2 = gr.Button("Predict on Built Molecule")
 
-    # Prediction Outputs
+    # Outputs
     with gr.Row():
         e_txt = gr.Text(label="Elastic")
         e_img = gr.Image(type="pil", label="Elastic Attention")