Update app.py

app.py

@@ -9,6 +9,7 @@ from huggingface_hub import hf_hub_download
 # ---------------------------
 # Device configuration
 # ---------------------------
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ---------------------------
@@ -16,7 +17,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # ---------------------------
 
 def build_model():
-    """Re‑create the network architecture used during training."""
+    """Re-create the network architecture used during training."""
     backbone = timm.create_model(
         "convnext_small", pretrained=False, num_classes=0, global_pool="avg"
     )
@@ -57,19 +58,29 @@ transform = transforms.Compose(
 # ---------------------------
 # Inference function
 # ---------------------------
-THRESHOLD = 0.5  # adjust if you want to tweak the decision boundary
 
-def predict(img: Image.Image, magnification: int, ra_conc: float):
-    """Return probabilities for High / Low CPM classes.
 
-    Note: `magnification` and `ra_conc` are recorded but *not* used for
-    inference. They are included so that users can tag inputs with these
-    metadata values.
+def predict(
+    img: Image.Image,
+    magnification: int,
+    ra_conc: float,
+    temperature: float = 1.0,
+):
+    """Return probabilities for High/Low CPM classes with optional temperature scaling.
+
+    Args:
+        img: Microscopy image.
+        magnification: Tag for objective magnification (×4/10/20).
+        ra_conc: Tag for RA concentration (µM).
+        temperature: Temperature parameter for confidence calibration. T>1 lowers
+            confidence, T<1 increases confidence.
     """
     img_tensor = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
         logit = model(img_tensor)
-        prob_high = torch.sigmoid(logit).item()
+        # Temperature scaling for calibration
+        logit_scaled = logit / temperature
+        prob_high = torch.sigmoid(logit_scaled).item()
 
     prob_low = 1.0 - prob_high
 
@@ -89,13 +100,21 @@ demo = gr.Interface(
         gr.Image(type="pil", label="Microscopy Image"),
         gr.Dropdown(choices=[4, 10, 20], value=10, label="Magnification (×)"),
         gr.Dropdown(choices=[0.1, 0.5, 1.0], value=0.1, label="RA Concentration (µM)"),
+        gr.Slider(
+            minimum=0.5,
+            maximum=5.0,
+            step=0.1,
+            value=1.0,
+            label="Temperature (confidence calibration)",
+            info="Increase temperature to reduce overconfidence",
+        ),
     ],
     outputs=gr.Label(num_top_classes=2, label="Predicted CPM Class & Probability"),
     title="iPS Cell Quality Classifier",
     description=(
-        "Upload a microscopy image, choose magnification and RA concentration "
-        "(for record‑keeping), and instantly get the predicted CPM quality "
-        "class with its probability."
+        "Upload a microscopy image, choose magnification & RA concentration "
+        "(metadata only), then optionally adjust the *Temperature* slider to "
+        "calibrate confidence if predictions look over‑ or under‑confident."
     ),
 )