Update

backend/app.py

@@ -509,6 +509,8 @@ async def predict(model_name: str = Form(...), file: UploadFile = File(...)):
                 content={"error": "CIN classifier not available on server."},
                 status_code=503,
             )
+
+        # Decode uploaded image and run colposcopy detector
         nparr = np.frombuffer(contents, np.uint8)
         img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
         results = yolo_colposcopy.predict(source=img, conf=0.7, save=False, verbose=False)
@@ -519,36 +521,52 @@ async def predict(model_name: str = Form(...), file: UploadFile = File(...)):
         crop = img[y1:y2, x1:x2]
         crop = cv2.resize(crop, (224, 224))
         img_t = transform(crop).unsqueeze(0).to(device)
+
+        # Extract features from multiple ResNet backbones
         f50 = extract_cbf_features(resnet50_blocks, img_t)
         f101 = extract_cbf_features(resnet101_blocks, img_t)
         f152 = extract_cbf_features(resnet152_blocks, img_t)
         features = np.concatenate([f50, f101, f152]).reshape(1, -1)
+
+        # Scale and predict
         X_scaled = MinMaxScaler().fit_transform(features)
-        pred = clf.predict(X_scaled)[0]
-        proba = clf.predict_proba(X_scaled)[0]
-        # Get actual number of classes from model output
-        classes = ["Low-grade", "High-grade"]  # Binary CIN classification
-        predicted_label = classes[pred]
-        confidences = {classes[i]: float(proba[i]) for i in range(len(classes))}
-        
-        # Map to more detailed classification based on confidence
-        if predicted_label == "High-grade" and confidences["High-grade"] > 0.8:
-            detailed_class = "CIN3"
-        elif predicted_label == "High-grade":
-            detailed_class = "CIN2"
-        else:
-            detailed_class = "CIN1"
 
-        # Average / primary confidence for display
+        # Ensure classifier supports probability outputs
+        try:
+            proba = clf.predict_proba(X_scaled)[0]
+        except Exception as e:
+            return JSONResponse(
+                content={"error": "CIN classifier does not support probability estimates (predict_proba)."},
+                status_code=503,
+            )
+
+        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:
+            return JSONResponse(
+                content={
+                    "error": "CIN classifier must output 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 a brief AI interpretation using the Mistral client (if available)
+        # Generate AI interpretation using Mistral client (if available)
         ai_interp = generate_cin_summary(predicted_label, confidences, avg_confidence)
 
         return {
             "model_used": "CIN Classifier",
-            "prediction": detailed_class,
-            "grade": predicted_label,
+            "prediction": predicted_label,
             "confidence": confidences,
             "summary": {
                 "avg_confidence": round(avg_confidence, 2),