Update app.py

app.py

@@ -1,14 +1,17 @@
-# We import gradio which is the library we use to build the web interface
+# gradio is the library used to build the web interface
 import gradio as gr
 
-# numpy is used for numerical operations and array manipulation
+# numpy is used for numerical operations
 import numpy as np
 
-# ai_edge_litert is Google's official replacement for tf.lite
-# it is lighter, faster, and works correctly on Python 3.13
-from ai_edge_litert.interpreter import Interpreter
+# torch is the core PyTorch library used to run the model
+import torch
 
-# PIL (Pillow) is used to handle image loading and resizing
+# torchvision provides the ResNet50 architecture and image transforms
+import torchvision.transforms as transforms
+from torchvision import models
+
+# PIL is used for image loading and conversion
 from PIL import Image
 
 
@@ -16,58 +19,48 @@ from PIL import Image
 # LOAD THE MODEL
 # ------------------------------------
 
-# This loads the tflite model file from the current directory
-# Interpreter is the ai_edge_litert equivalent of tf.lite.Interpreter
-interpreter = Interpreter(model_path="best_gatekeeper_v2.pth")
-
-# This allocates memory for the model's input and output tensors
-# You must always call this before running inference
-interpreter.allocate_tensors()
-
-# This gets the details of the input tensor
-# It tells us the expected shape, data type, and index of the input
-input_details = interpreter.get_input_details()
-
-# This gets the details of the output tensor
-# It tells us the shape and index of the output so we can read predictions
-output_details = interpreter.get_output_details()
+# 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}")
 
-# This is the image size the model expects
-# ResNet50 was trained on 224x224 images so we keep this the same
-INPUT_SIZE = (224, 224)
+# 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:
+# class 0 = Non-Cervix, class 1 = Cervix
+model.fc = torch.nn.Linear(model.fc.in_features, 2)
 
-# ------------------------------------
-# IMAGE PREPROCESSING FUNCTION
-# ------------------------------------
+# we load the saved weights from the .pth file
+# map_location=device ensures it loads correctly even without a GPU
+state_dict = torch.load("best_gatekeeper_v2.pth", map_location=device)
+model.load_state_dict(state_dict)
 
-def preprocess_image(image):
-    # image comes in as a numpy array from gradio
-    # we convert it to a PIL Image object so we can resize it easily
-    # we also make sure it is in RGB format (3 channels: red, green, blue)
-    img = Image.fromarray(image).convert("RGB")
+# we move the model to the correct device (CPU or GPU)
+model = model.to(device)
 
-    # we resize the image to match what the model expects
-    # if the image is not 224x224 the model will throw a shape error
-    img = img.resize(INPUT_SIZE)
+# we set the model to evaluation mode
+# this disables dropout and batch normalisation training behaviour
+model.eval()
 
-    # we convert the PIL image back to a numpy array
-    # dtype=np.float32 is important because the model expects 32-bit floats
-    img = np.array(img, dtype=np.float32)
+print("Gatekeeper model loaded successfully")
 
-    # we divide all pixel values by 255
-    # this converts pixel values from the range [0, 255] to [0, 1]
-    # this is called normalization and it helps the model perform correctly
-    img = img / 255.0
+# this is the image size ResNet50 expects
+INPUT_SIZE = 224
 
-    # the model expects a batch of images, not a single image
-    # so we add an extra dimension at position 0
-    # this changes the shape from (224, 224, 3) to (1, 224, 224, 3)
-    # the 1 represents a batch size of 1 (one image at a time)
-    img = np.expand_dims(img, axis=0)
+# these are the standard ImageNet normalisation values
+# ResNet50 was pretrained on ImageNet so we use the same values
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD  = [0.229, 0.224, 0.225]
 
-    # we return the fully preprocessed image ready for inference
-    return img
+# we define the preprocessing pipeline using torchvision transforms
+preprocess = transforms.Compose([
+    transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
+    transforms.ToTensor(),                          # converts [0,255] → [0,1]
+    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
+])
 
 
 # ------------------------------------
@@ -75,53 +68,40 @@ def preprocess_image(image):
 # ------------------------------------
 
 def classify_image(image):
-    # if the user clicks the button without uploading an image
-    # we return None for the scores and a warning message
+    # if the user submits without an image return a warning
     if image is None:
         return None, "Please upload an image first"
 
-    # we send the image through our preprocessing function
-    processed = preprocess_image(image)
-
-    # we load the preprocessed image into the model's input tensor
-    # input_details[0]['index'] gives us the correct tensor index to write to
-    interpreter.set_tensor(input_details[0]['index'], processed)
-
-    # this actually runs the model on the input we just loaded
-    interpreter.invoke()
+    # convert the numpy array from gradio to a PIL Image in RGB format
+    img = Image.fromarray(image).convert("RGB")
 
-    # this reads the output from the model after inference is complete
-    # output_details[0]['index'] gives us the correct tensor index to read from
-    output = interpreter.get_tensor(output_details[0]['index'])
+    # apply the preprocessing pipeline and add a batch dimension
+    # unsqueeze(0) changes shape from (3, 224, 224) to (1, 3, 224, 224)
+    tensor = preprocess(img).unsqueeze(0).to(device)
 
-    # we print the raw output to the console for debugging purposes
-    # this is useful to confirm the model is producing expected values
-    print(f"Raw model output: {output}")
+    # run inference without computing gradients (saves memory and is faster)
+    with torch.no_grad():
+        output = model(tensor)          # raw logits shape: (1, 2)
+        probs = torch.softmax(output, dim=1)[0]  # convert to probabilities
 
-    # index 0 of the output corresponds to Non-Cervix probability
-    # we convert it to a plain Python float for easier handling
-    prob_non_cervix = float(output[0][0])
+    # extract individual class probabilities as plain Python floats
+    prob_non_cervix = float(probs[0])
+    prob_cervix     = float(probs[1])
 
-    # index 1 of the output corresponds to Cervix probability
-    prob_cervix = float(output[0][1])
+    print(f"Non-Cervix: {prob_non_cervix:.4f} | Cervix: {prob_cervix:.4f}")
 
-    # we compare the two probabilities to determine the final prediction
-    # whichever class has the higher probability is our prediction
+    # determine the final prediction label
     if prob_cervix > prob_non_cervix:
         prediction_text = "Cervix Detected"
     else:
         prediction_text = "Non-Cervix"
 
-    # we build a dictionary of class names mapped to their confidence scores
-    # gradio's Label component accepts this format and displays it as a bar chart
-    # we round to 4 decimal places to keep the display clean
+    # 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)
+        "Non-Cervix": round(prob_non_cervix, 4),
     }
 
-    # we return both the scores dictionary and the prediction text
-    # these map to the two output components in the gradio interface
     return scores, prediction_text
 
 
@@ -129,11 +109,8 @@ def classify_image(image):
 # GRADIO USER INTERFACE
 # ------------------------------------
 
-# gr.Blocks gives us full control over the layout of the interface
-# theme=gr.themes.Soft() gives it a clean and soft visual style
 with gr.Blocks(theme=gr.themes.Soft()) as app:
 
-    # gr.Markdown renders formatted text at the top of the page
     gr.Markdown("""
     # Gatekeeper Model
     ### Cervix Image Binary Classifier
@@ -141,58 +118,35 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
     ---
     """)
 
-    # gr.Row arranges the components inside it horizontally side by side
     with gr.Row():
 
-        # the first column holds the input components on the left side
         with gr.Column():
-
-            # gr.Image creates an image upload box
-            # type="numpy" means the image will be passed to our function
-            # as a numpy array which is what we need for preprocessing
             input_image = gr.Image(
                 label="Upload Image",
                 type="numpy"
             )
-
-            # this is the main button the user clicks to run the model
-            # variant="primary" makes it stand out visually as the main action
-            # size="lg" makes it large and easy to click
             classify_btn = gr.Button(
                 "Run Classification",
                 variant="primary",
                 size="lg"
             )
-
-            # this is a secondary button to reset the interface
-            # variant="secondary" gives it a less prominent visual style
             clear_btn = gr.Button(
                 "Clear",
                 variant="secondary",
                 size="sm"
             )
 
-        # the second column holds the output components on the right side
         with gr.Column():
-
-            # gr.Label displays the confidence scores as a visual bar chart
-            # num_top_classes=2 tells it to show both classes
             output_scores = gr.Label(
                 label="Confidence Scores",
                 num_top_classes=2
             )
-
-            # gr.Textbox displays the final prediction as plain text
-            # interactive=False means the user cannot edit it
-            # it is read-only output only
             output_text = gr.Textbox(
                 label="Prediction",
                 interactive=False,
                 text_align="center"
             )
 
-    # this adds a reference table at the bottom so users understand
-    # what the two class indices mean
     gr.Markdown("""
     ---
     | Index | Label       | Meaning                          |
@@ -205,29 +159,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
     It is not intended for clinical diagnosis or medical use.
     """)
 
-    # ------------------------------------
-    # BUTTON ACTIONS
-    # ------------------------------------
-
-    # this connects the classify button to the classify_image function
-    # inputs tells gradio which component to read from
-    # outputs tells gradio which components to write the results to
     classify_btn.click(
         fn=classify_image,
         inputs=input_image,
         outputs=[output_scores, output_text]
     )
 
-    # this connects the clear button to a simple lambda function
-    # a lambda is a small anonymous function defined in one line
-    # it returns None for the image, None for scores, and empty string for text
-    # this effectively resets all three components back to their empty state
     clear_btn.click(
         fn=lambda: (None, None, ""),
         inputs=None,
         outputs=[input_image, output_scores, output_text]
     )
 
-# this starts the gradio web server and launches the interface
-# on hugging face spaces this is called automatically
 app.launch()