Reject non-leaf uploads in predictions

api.py

@@ -134,8 +134,31 @@ def load_model():
 def preprocess_image(image_bytes: bytes) -> np.ndarray:
     image = Image.open(BytesIO(image_bytes)).convert("RGB")
     image = image.resize((224, 224))
-    img_array = np.array(image, dtype=np.float32) / 255.0
-    return np.expand_dims(img_array, axis=0)
+    rgb_uint8 = np.array(image, dtype=np.uint8)
+    img_array = rgb_uint8.astype(np.float32) / 255.0
+    return np.expand_dims(img_array, axis=0), rgb_uint8
+
+
+def should_reject_as_no_leaf(rgb_uint8: np.ndarray, probs: np.ndarray) -> tuple[bool, str]:
+    confidence = float(np.max(probs))
+    sorted_probs = np.sort(probs)
+    second_best = float(sorted_probs[-2]) if sorted_probs.size > 1 else 0.0
+    margin = confidence - second_best
+
+    r = rgb_uint8[:, :, 0].astype(np.float32)
+    g = rgb_uint8[:, :, 1].astype(np.float32)
+    b = rgb_uint8[:, :, 2].astype(np.float32)
+
+    green_mask = (g > 40) & (g > 1.05 * r) & (g > 1.05 * b)
+    green_ratio = float(np.mean(green_mask))
+
+    if confidence < 0.55:
+        return True, "low_confidence"
+    if margin < 0.12 and confidence < 0.70:
+        return True, "ambiguous_prediction"
+    if green_ratio < 0.05 and confidence < 0.80:
+        return True, "non_leaf_visual_pattern"
+    return False, "accepted"
 
 
 MODEL = load_model()
@@ -154,7 +177,7 @@ async def predict(file: UploadFile = File(...)) -> dict:
 
     try:
         image_bytes = await file.read()
-        input_tensor = preprocess_image(image_bytes)
+        input_tensor, rgb_uint8 = preprocess_image(image_bytes)
         preds = MODEL.predict(input_tensor)
         probs = np.asarray(preds[0], dtype=np.float32)
 
@@ -176,6 +199,19 @@ async def predict(file: UploadFile = File(...)) -> dict:
                 "confidence": float(probs[idx]),
             })
 
+        reject, reject_reason = should_reject_as_no_leaf(rgb_uint8, probs)
+        if reject:
+            return {
+                "prediction": {
+                    "class_index": -1,
+                    "class_name": "No leaf detected",
+                    "confidence": float(1.0 - confidence),
+                },
+                "top3": top3,
+                "rejected": True,
+                "reject_reason": reject_reason,
+            }
+
         return {
             "prediction": {
                 "class_index": result_idx,
@@ -183,6 +219,7 @@ async def predict(file: UploadFile = File(...)) -> dict:
                 "confidence": confidence,
             },
             "top3": top3,
+            "rejected": False,
         }
     except HTTPException:
         raise