Update app.py

app.py

@@ -64,4 +64,60 @@ except AttributeError:
     label_names = ["Background", "Caries", "Normal Teeth", "Plaque"]
 
 def preprocess_image(image):
-    """Load and preprocess a
+    """Load and preprocess a PIL image."""
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    image = image.convert("RGB")
+    return inference_transform(image).unsqueeze(0)
+
+def predict_image(image):
+    """Run inference on image and return multi-label predictions."""
+    pixel_values = preprocess_image(image)
+    
+    with torch.no_grad():
+        logits = quantized_model(pixel_values)
+    
+    probs = torch.sigmoid(logits).squeeze(0)
+    preds = (probs > 0.5).int().tolist()
+    
+    detected_conditions = []
+    for i, (label, pred) in enumerate(zip(label_names, preds)):
+        if pred == 1:
+            confidence = probs[i].item()
+            detected_conditions.append(f"{label} (confidence: {confidence:.2%})")
+    
+    # Check for potential Caries
+    try:
+        caries_index = label_names.index("Caries")
+        caries_prob = probs[caries_index].item()
+        if 0.3 <= caries_prob < 0.5:
+            detected_conditions.append(f"Possible Caries (confidence: {caries_prob:.2%})")
+    except ValueError:
+        pass
+    
+    if detected_conditions:
+        result = "Detected: " + ", ".join(detected_conditions)
+    else:
+        result = "No dental issues detected"
+    
+    return result
+
+# Example images
+examples = [
+    ["example_image1.jfif"],
+    ["example_image2.jfif"],
+    ["example_image3.jfif"]
+]
+
+# Gradio interface
+iface = gr.Interface(
+    fn=predict_image,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Dental Image Multi-Label Classification",
+    description="Upload an image or select from the examples below to predict dental conditions. The model can detect multiple dental issues in a single image.",
+    examples=examples
+)
+
+if __name__ == "__main__":
+    iface.launch()