Deploy KidneyDL CT Scan Classifier

src/cnnClassifier/pipeline/prediction.py

@@ -3,9 +3,19 @@ from tensorflow.keras.models import load_model
 from tensorflow.keras.preprocessing import image
 import os
 
-# Minimum softmax confidence required to trust a prediction.
-# Below this threshold the image is likely not a kidney CT scan.
-CONFIDENCE_THRESHOLD = 0.80
+# CT scans are grayscale — all three RGB channels are nearly identical.
+# This threshold is the max allowed std-dev of inter-channel pixel differences.
+# True CT scans score ~0–8; colour photos score ~20–80+.
+GRAYSCALE_THRESHOLD = 15.0
+
+
+def _is_ct_like(img_array):
+    """Return True if the image looks like a grayscale CT scan."""
+    r = img_array[:, :, 0].astype(float)
+    g = img_array[:, :, 1].astype(float)
+    b = img_array[:, :, 2].astype(float)
+    max_channel_diff = max(np.std(r - g), np.std(r - b), np.std(g - b))
+    return max_channel_diff < GRAYSCALE_THRESHOLD
 
 
 class PredictionPipeline:
@@ -24,15 +34,16 @@ class PredictionPipeline:
             model = load_model(model_path)
 
         img = image.load_img(self.filename, target_size=(224, 224))
-        img_array = image.img_to_array(img)
-        img_array = np.expand_dims(img_array, axis=0) / 255.0
+        img_array = image.img_to_array(img)   # shape (224, 224, 3), values 0–255
 
-        predictions = model.predict(img_array)
-        confidence  = float(np.max(predictions))
-        class_idx   = int(np.argmax(predictions, axis=1)[0])
+        # Reject non-CT images before they reach the model
+        if not _is_ct_like(img_array):
+            return [{"image": "InvalidImage"}]
 
-        if confidence < CONFIDENCE_THRESHOLD:
-            return [{"image": "InvalidImage", "confidence": round(confidence, 4)}]
+        img_input = np.expand_dims(img_array, axis=0) / 255.0
+        predictions = model.predict(img_input)
+        class_idx   = int(np.argmax(predictions, axis=1)[0])
+        confidence  = float(np.max(predictions))
 
         return [{"image": "Tumor" if class_idx == 1 else "Normal",
                  "confidence": round(confidence, 4)}]