Update app.py

app.py

@@ -20,48 +20,57 @@ from PIL import Image
 # ------------------------------------
 
 # we detect whether a GPU is available and fall back to CPU if not
-# hugging face free tier runs on CPU so this will almost always be cpu
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Running on: {device}")
 
 # we recreate the ResNet50 architecture
-# weights=None because we will load our own trained weights below
 model = models.resnet50(weights=None)
 
-# the original ResNet50 outputs 1000 classes (ImageNet)
-# we replace the final fully connected layer to output 2 classes:
+# replace the final fully connected layer to output 2 classes:
 # class 0 = Non-Cervix, class 1 = Cervix
 model.fc = torch.nn.Linear(model.fc.in_features, 2)
 
-# we load the saved weights from the .pth file
-# map_location=device ensures it loads correctly even without a GPU
+# load the saved weights from the .pth file
 state_dict = torch.load("best_gatekeeper_v2.pth", map_location=device)
 model.load_state_dict(state_dict)
 
-# we move the model to the correct device (CPU or GPU)
+# move the model to the correct device and set to evaluation mode
 model = model.to(device)
-
-# we set the model to evaluation mode
-# this disables dropout and batch normalisation training behaviour
 model.eval()
 
 print("Gatekeeper model loaded successfully")
 
-# this is the image size ResNet50 expects
+# image size ResNet50 expects
 INPUT_SIZE = 224
 
-# these are the standard ImageNet normalisation values
-# ResNet50 was pretrained on ImageNet so we use the same values
+# standard ImageNet normalisation values used during pretraining
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD  = [0.229, 0.224, 0.225]
 
-# we define the preprocessing pipeline using torchvision transforms
+# preprocessing pipeline
 preprocess = transforms.Compose([
     transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
-    transforms.ToTensor(),                          # converts [0,255] → [0,1]
+    transforms.ToTensor(),
     transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
 ])
 
+# ------------------------------------
+# THRESHOLDS
+# ------------------------------------
+
+# minimum probability for cervix to be accepted as a positive detection
+CERVIX_THRESHOLD = 0.55
+
+# minimum gap between cervix and non-cervix probabilities
+# if the gap is smaller than this the prediction is too uncertain to trust
+CONFIDENCE_GAP = 0.15
+
+# minimum image brightness - images below this are too dark to classify
+MIN_BRIGHTNESS = 30
+
+# minimum image contrast - images below this are blank or uniform
+MIN_STD = 20
+
 
 # ------------------------------------
 # CLASSIFICATION FUNCTION
@@ -72,32 +81,43 @@ def classify_image(image):
     if image is None:
         return None, "Please upload an image first"
 
-    # convert the numpy array from gradio to a PIL Image in RGB format
-    img = Image.fromarray(image).convert("RGB")
+    # Option 3: Basic image sanity checks
+    # run these before the model to catch obviously bad images early
+    img_array = np.array(image)
 
-    # apply the preprocessing pipeline and add a batch dimension
-    # unsqueeze(0) changes shape from (3, 224, 224) to (1, 3, 224, 224)
+    # reject images that are too dark to analyse reliably
+    if img_array.mean() < MIN_BRIGHTNESS:
+        return None, "Image is too dark - please upload a clearer photo"
+
+    # reject images that are blank, uniformly coloured, or plain screenshots
+    if img_array.std() < MIN_STD:
+        return None, "Image appears blank or uniform - please upload a real photo"
+
+    # Preprocess
+    img = Image.fromarray(image).convert("RGB")
     tensor = preprocess(img).unsqueeze(0).to(device)
 
-    # run inference without computing gradients (saves memory and is faster)
+    # Run inference
     with torch.no_grad():
-        output = model(tensor)          # raw logits shape: (1, 2)
-        probs = torch.softmax(output, dim=1)[0]  # convert to probabilities
+        output = model(tensor)
+        probs = torch.softmax(output, dim=1)[0]
 
-    # extract individual class probabilities as plain Python floats
     prob_non_cervix = float(probs[0])
     prob_cervix     = float(probs[1])
 
     print(f"Non-Cervix: {prob_non_cervix:.4f} | Cervix: {prob_cervix:.4f}")
 
-    # determine the final prediction label
-# cervix must score at least 0.55 to be accepted as a positive detection
-    if prob_cervix >= 0.55:
+    # Option 1: Confidence threshold + gap check
+    gap = prob_cervix - prob_non_cervix
+
+    if prob_cervix >= CERVIX_THRESHOLD and gap >= CONFIDENCE_GAP:
         prediction_text = "Cervix Detected"
-    else:
+    elif prob_non_cervix >= CERVIX_THRESHOLD and gap <= -CONFIDENCE_GAP:
         prediction_text = "Non-Cervix"
+    else:
+        # not confident enough either way - temporary misclassification safety net
+        prediction_text = "Uncertain - please retake or upload a clearer image"
 
-    # build a dictionary for gradio's Label component (displays as bar chart)
     scores = {
         "Cervix":     round(prob_cervix, 4),
         "Non-Cervix": round(prob_non_cervix, 4),
@@ -155,6 +175,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
     | 0     | Non-Cervix  | Image does NOT contain cervix    |
     | 1     | Cervix      | Image contains cervix            |
 
+    ---
+    **How predictions work:**
+    - **Cervix Detected** - model scored >= 0.55 with a gap of >= 0.15 over Non-Cervix
+    - **Non-Cervix** - model scored >= 0.55 with a gap of >= 0.15 over Cervix
+    - **Uncertain** - model was not confident enough; retake the image
+
     ---
     Disclaimer: This tool is for research purposes only.
     It is not intended for clinical diagnosis or medical use.