Update app.py

app.py

@@ -11,43 +11,43 @@ interpreter.allocate_tensors()
 input_details = interpreter.get_input_details()
 output_details = interpreter.get_output_details()
 
+# Class labels
 class_names = ['Dyskeratotic', 'Koilocytotic', 'Metaplastic', 'Parabasal', 'Superficial-Intermediat']
 CONFIDENCE_THRESHOLD = 0.7
+TEMPERATURE = 0.5  # 🔑 Sharpening factor (lower = more confident)
 
 def preprocess_image(image):
-    """Preprocess input image for model"""
+    """Resize & normalize image for float32 TFLite model"""
     image = image.resize((224, 224)).convert("RGB")
-    img_array = np.array(image, dtype=np.float32)
-
-    # 🔑 IMPORTANT: Check model input scale
-    if input_details[0]['dtype'] == np.uint8:
-        img_array = np.expand_dims(img_array, axis=0).astype(np.uint8)
-    else:
-        img_array = img_array / 255.0
-        img_array = np.expand_dims(img_array, axis=0).astype(np.float32)
-
+    img_array = np.array(image, dtype=np.float32) / 255.0
+    img_array = np.expand_dims(img_array, axis=0)
     return img_array
 
 def predict_image(image):
     try:
-        # Preprocess
         img_array = preprocess_image(image)
 
-        # Inference
+        # Run inference
         interpreter.set_tensor(input_details[0]['index'], img_array)
         interpreter.invoke()
         output = interpreter.get_tensor(output_details[0]['index'])
 
-        # Ensure output is softmax
-        probabilities = tf.nn.softmax(output[0]).numpy()
+        # Apply softmax with temperature scaling (sharper confidence)
+        probabilities = tf.nn.softmax(output[0] / TEMPERATURE).numpy()
 
+        # Top prediction
         class_idx = int(np.argmax(probabilities))
         confidence = float(np.max(probabilities))
 
-        if confidence < CONFIDENCE_THRESHOLD:
-            return f"⚠️ Low confidence ({confidence:.2f}). Try another clearer image."
-        else:
+        if confidence >= CONFIDENCE_THRESHOLD:
             return f"✅ Prediction: {class_names[class_idx]} (Confidence: {confidence:.2f})"
+        else:
+            # Show Top-2 predictions if low confidence
+            top2_idx = np.argsort(probabilities)[-2:][::-1]
+            suggestion = ", ".join(
+                [f"{class_names[i]} ({probabilities[i]:.2f})" for i in top2_idx]
+            )
+            return f"⚠️ Low confidence ({confidence:.2f}).\nPossible classes: {suggestion}"
 
     except Exception as e:
         return f"❌ Error: {str(e)}"
@@ -58,5 +58,5 @@ gr.Interface(
     inputs=gr.Image(type="pil", label="Upload Cervical Cell Image"),
     outputs="text",
     title="Cervical Cancer Cell Classification",
-    description="Upload an image to classify cervical cells using a TFLite model."
+    description="Upload a cervical cell image. The model predicts the cell type with confidence scores."
 ).launch()