Update app.py

app.py

@@ -4,12 +4,8 @@ import gradio as gr
 # numpy is used for numerical operations
 import numpy as np
 
-# torch is the core PyTorch library used to run the model
-import torch
-
-# torchvision provides the ResNet50 architecture and image transforms
-import torchvision.transforms as transforms
-from torchvision import models
+# ai_edge_litert is Google's official TFLite runtime
+from ai_edge_litert.interpreter import Interpreter
 
 # PIL is used for image loading and conversion
 from PIL import Image
@@ -19,40 +15,21 @@ from PIL import Image
 # LOAD THE MODEL
 # ------------------------------------
 
-# we detect whether a GPU is available and fall back to CPU if not
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Running on: {device}")
-
-# we recreate the ResNet50 architecture
-model = models.resnet50(weights=None)
-
-# 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)
-
-# 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)
+# load the float32 TFLite model
+interpreter = Interpreter(model_path="resnet50_float32.tflite")
 
-# move the model to the correct device and set to evaluation mode
-model = model.to(device)
-model.eval()
+# allocate memory for the model's input and output tensors
+interpreter.allocate_tensors()
 
-print("Gatekeeper model loaded successfully")
+# get input and output tensor details
+input_details  = interpreter.get_input_details()
+output_details = interpreter.get_output_details()
 
 # image size ResNet50 expects
-INPUT_SIZE = 224
+INPUT_SIZE = (224, 224)
 
-# standard ImageNet normalisation values used during pretraining
-IMAGENET_MEAN = [0.485, 0.456, 0.406]
-IMAGENET_STD  = [0.229, 0.224, 0.225]
+print("Gatekeeper model loaded successfully")
 
-# preprocessing pipeline
-preprocess = transforms.Compose([
-    transforms.Resize((INPUT_SIZE, INPUT_SIZE)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
-])
 
 # ------------------------------------
 # THRESHOLDS
@@ -72,6 +49,23 @@ MIN_BRIGHTNESS = 30
 MIN_STD = 20
 
 
+# ------------------------------------
+# IMAGE PREPROCESSING FUNCTION
+# ------------------------------------
+
+def preprocess_image(image):
+    # convert numpy array to PIL Image in RGB format and resize
+    img = Image.fromarray(image).convert("RGB").resize(INPUT_SIZE)
+
+    # convert to float32 numpy array and normalise to [0, 1]
+    img = np.array(img, dtype=np.float32) / 255.0
+
+    # add batch dimension: (224, 224, 3) → (1, 224, 224, 3)
+    img = np.expand_dims(img, axis=0)
+
+    return img
+
+
 # ------------------------------------
 # CLASSIFICATION FUNCTION
 # ------------------------------------
@@ -82,7 +76,6 @@ def classify_image(image):
         return None, "Please upload an image first"
 
     # Option 3: Basic image sanity checks
-    # run these before the model to catch obviously bad images early
     img_array = np.array(image)
 
     # reject images that are too dark to analyse reliably
@@ -93,17 +86,22 @@ def classify_image(image):
     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)
+    # preprocess the image
+    processed = preprocess_image(image)
+
+    # load the preprocessed image into the model's input tensor
+    interpreter.set_tensor(input_details[0]['index'], processed)
+
+    # run inference
+    interpreter.invoke()
 
-    # Run inference
-    with torch.no_grad():
-        output = model(tensor)
-        probs = torch.softmax(output, dim=1)[0]
+    # read the output tensor
+    output = interpreter.get_tensor(output_details[0]['index'])
+    print(f"Raw model output: {output}")
 
-    prob_non_cervix = float(probs[0])
-    prob_cervix     = float(probs[1])
+    # extract individual class probabilities
+    prob_non_cervix = float(output[0][0])
+    prob_cervix     = float(output[0][1])
 
     print(f"Non-Cervix: {prob_non_cervix:.4f} | Cervix: {prob_cervix:.4f}")