Update

backend/app.py

@@ -542,29 +542,56 @@ async def predict(model_name: str = Form(...), file: UploadFile = File(...)):
 
         num_classes = int(len(proba))
 
-        # Expecting a 3-class CIN classifier. If not 3, return a clear error so the user can supply a 3-class model.
-        if num_classes != 3:
+        # Handle different classifier output sizes:
+        # - If 3 classes: map directly to CIN1/CIN2/CIN3
+        # - If 2 classes: apply a conservative heuristic to split High-grade into CIN2/CIN3
+        if num_classes == 3:
+            classes = ["CIN1", "CIN2", "CIN3"]
+            confidences = {classes[i]: float(proba[i]) for i in range(3)}
+            predicted_idx = int(np.argmax(proba))
+            predicted_label = classes[predicted_idx]
+            avg_confidence = float(np.max(proba)) * 100
+            mapping_used = "direct_3class"
+        elif num_classes == 2:
+            # Binary model detected (e.g., Low-grade vs High-grade). We'll convert to CIN1/CIN2/CIN3
+            # Heuristic:
+            # - CIN1 <- low_grade_prob
+            # - Split high_grade_prob into CIN2 and CIN3 based on how confident 'high' is.
+            #   * If high <= 0.6 -> mostly CIN2
+            #   * If high >= 0.8 -> mostly CIN3
+            #   * Between 0.6 and 0.8 -> interpolate
+            low_prob = float(proba[0])
+            high_prob = float(proba[1])
+
+            if high_prob <= 0.6:
+                cin3_factor = 0.0
+            elif high_prob >= 0.8:
+                cin3_factor = 1.0
+            else:
+                cin3_factor = (high_prob - 0.6) / 0.2
+
+            cin1 = low_prob
+            cin3 = high_prob * cin3_factor
+            cin2 = high_prob - cin3
+
+            confidences = {"CIN1": cin1, "CIN2": cin2, "CIN3": cin3}
+            # pick highest of the mapped three as primary prediction
+            predicted_label = max(confidences.items(), key=lambda x: x[1])[0]
+            avg_confidence = float(max(confidences.values())) * 100
+            mapping_used = "binary_to_3class_heuristic"
+        else:
             return JSONResponse(
                 content={
-                    "error": "CIN classifier must output 3-class probabilities (CIN1, CIN2, CIN3).",
+                    "error": "CIN classifier must output 2-class (Low/High) or 3-class probabilities (CIN1, CIN2, CIN3).",
                     "detected_num_classes": num_classes,
                 },
                 status_code=503,
             )
 
-        # Map probabilities to explicit CIN labels (CIN1, CIN2, CIN3)
-        classes = ["CIN1", "CIN2", "CIN3"]
-        confidences = {classes[i]: float(proba[i]) for i in range(3)}
-
-        # Primary prediction is the class with the highest probability
-        predicted_idx = int(np.argmax(proba))
-        predicted_label = classes[predicted_idx]
-        avg_confidence = float(np.max(proba)) * 100
-
         # Generate AI interpretation using Mistral client (if available)
         ai_interp = generate_cin_summary(predicted_label, confidences, avg_confidence)
 
-        return {
+        response = {
             "model_used": "CIN Classifier",
             "prediction": predicted_label,
             "confidence": confidences,
@@ -573,6 +600,16 @@ async def predict(model_name: str = Form(...), file: UploadFile = File(...)):
                 "ai_interpretation": ai_interp,
             },
         }
+
+        # If we used the binary->3class heuristic, include a diagnostic field so callers know it was mapped
+        if 'mapping_used' in locals() and mapping_used == 'binary_to_3class_heuristic':
+            response["mapping_used"] = mapping_used
+            response["mapping_note"] = (
+                "The server mapped a binary Low/High classifier to CIN1/CIN2/CIN3 using a heuristic split. "
+                "This is an approximation — for clinical use please supply a native 3-class model."
+            )
+
+        return response
     elif model_name == "histopathology":
             result = predict_histopathology(image)
             return result