Reject uncertain traffic sign predictions

app.py

@@ -26,6 +26,10 @@ DATABASE_PATH = BASE_DIR / "instance" / "traffic_signs.sqlite3"
 MODEL_PATH = BASE_DIR / "traffic_classifier.h5"
 ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "webp"}
 NORMALIZE_INPUT = os.environ.get("NORMALIZE_INPUT", "false").lower() in {"1", "true", "yes"}
+MC_DROPOUT_PASSES = int(os.environ.get("MC_DROPOUT_PASSES", "12"))
+UNKNOWN_CONFIDENCE_THRESHOLD = float(os.environ.get("UNKNOWN_CONFIDENCE_THRESHOLD", "0.50"))
+UNKNOWN_MARGIN_THRESHOLD = float(os.environ.get("UNKNOWN_MARGIN_THRESHOLD", "0.10"))
+UNKNOWN_VOTE_THRESHOLD = float(os.environ.get("UNKNOWN_VOTE_THRESHOLD", "0.60"))
 URL_CONTENT_TYPES = {
     "image/jpeg": "jpg",
     "image/png": "png",
@@ -241,19 +245,55 @@ def save_uploaded_image(file):
 
 def prepare_image(path):
     image = Image.open(path).convert("RGB")
-    #image = image.resize((30, 30))
+    image = image.resize((30, 30))
     array = np.asarray(image, dtype=np.float32)
     if NORMALIZE_INPUT:
         array = array / 255.0
     return np.expand_dims(array, axis=0)
 
 
+def normalize_model_output(predictions):
+    predictions = np.asarray(predictions, dtype=np.float64)
+    total = float(np.sum(predictions))
+    if np.any(predictions < 0) or not np.isclose(total, 1.0, atol=1e-3):
+        shifted = predictions - np.max(predictions)
+        exp_values = np.exp(shifted)
+        return exp_values / np.sum(exp_values)
+    return predictions
+
+
+def predict_probabilities(batch):
+    if MC_DROPOUT_PASSES <= 1:
+        return normalize_model_output(model.predict(batch, verbose=0)[0]), 1.0
+
+    samples = []
+    votes = []
+    for _ in range(MC_DROPOUT_PASSES):
+        probabilities = normalize_model_output(model(batch, training=True).numpy()[0])
+        samples.append(probabilities)
+        votes.append(int(np.argmax(probabilities)))
+
+    mean_probabilities = np.mean(np.asarray(samples), axis=0)
+    vote_counts = np.bincount(votes, minlength=len(mean_probabilities))
+    vote_agreement = float(np.max(vote_counts) / MC_DROPOUT_PASSES)
+    return mean_probabilities, vote_agreement
+
+
 def predict_sign(path):
     if model is None:
         return "Model unavailable", 0.0
-    predictions = model.predict(prepare_image(path), verbose=0)[0]
-    class_index = int(np.argmax(predictions))
-    confidence = float(np.max(predictions))
+    probabilities, vote_agreement = predict_probabilities(prepare_image(path))
+    ranked_indices = np.argsort(probabilities)[::-1]
+    class_index = int(ranked_indices[0])
+    confidence = float(probabilities[class_index])
+    second_confidence = float(probabilities[int(ranked_indices[1])]) if len(ranked_indices) > 1 else 0.0
+    margin = confidence - second_confidence
+    if (
+        confidence < UNKNOWN_CONFIDENCE_THRESHOLD
+        or margin < UNKNOWN_MARGIN_THRESHOLD
+        or vote_agreement < UNKNOWN_VOTE_THRESHOLD
+    ):
+        return "Unknown traffic sign", confidence
     label = TRAFFIC_SIGN_CLASSES[class_index] if class_index < len(TRAFFIC_SIGN_CLASSES) else f"Class {class_index}"
     return label, confidence