Update app.py

app.py

@@ -1,71 +1,102 @@
+from huggingface_hub import hf_hub_download
 import torch
-from torchvision import models, transforms
+from transformers import ViTImageProcessor
 from PIL import Image
 import gradio as gr
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 
-# Updated class names with 'plaque' in front of 'calculus' and 'gingivitis'
-class_names = [
-    "Plaque Calculus",
-    "Caries",
-    "Plaque Gingivitis",
-    "Hypodontia",
-    "Mouth Ulcer",
-    "Tooth Discoloration"
-]
-
-# Load the model and update the final fully connected layer
-model = models.resnet50(weights=None)
-model.fc = torch.nn.Linear(model.fc.in_features, len(class_names))
-
-# Load the model weights from tooth_model.pth
-model.load_state_dict(torch.load('tooth_model.pth', map_location=torch.device('cpu')))
-model.eval()
+# Repository configuration
+REPO_ID = "IFMedTech/Dental_Q"
+MODEL_FILENAME = "quantized_model.ptl"
 
-# Preprocessing steps for input images
-preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-])
+# Download the model file from Hugging Face Hub
+model_path = hf_hub_download(
+    repo_id=REPO_ID,
+    filename=MODEL_FILENAME
+)
 
-def predict_image(image):
-    # Preprocess the image and add a batch dimension
-    processed_image = preprocess(image).unsqueeze(0)
+# Download the processor files (assuming they're in the same repo)
+processor_path = hf_hub_download(
+    repo_id=REPO_ID,
+    filename="preprocessor_config.json"
+)
 
-    with torch.no_grad():
-        outputs = model(processed_image)
-        probabilities = torch.nn.functional.softmax(outputs, dim=1)
-        top_probs, top_indices = torch.topk(probabilities, 2)  # Get top 2 predictions
+# Load Processor & Quantized Model
+processor = ViTImageProcessor.from_pretrained(REPO_ID)
+quantized_model = torch.jit.load(model_path, map_location="cpu")
+quantized_model.eval()
 
-        top_class_1 = class_names[top_indices[0][0]]
-        top_prob_1 = top_probs[0][0].item()
+# Define Inference Preprocessing
+size = processor.size['height']
+normalize = Normalize(mean=processor.image_mean, std=processor.image_std)
+inference_transform = Compose([
+    Resize(size),
+    CenterCrop(size),
+    ToTensor(),
+    normalize
+])
 
-        # Initialize result with the top prediction
-        result = top_class_1
+# Multi-label class names
+try:
+    label_names = [quantized_model.config.id2label[i] for i in range(len(quantized_model.config.id2label))]
+except AttributeError:
+    label_names = ["Background", "Caries", "Normal Teeth", "Plaque"]
 
-        # Include the second prediction if the top prediction's probability is less than 80%
-        if top_prob_1 < 0.8:
-            top_class_2 = class_names[top_indices[0][1]]
-            result += f", {top_class_2}"
+def preprocess_image(image):
+    """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 to use as input
+# Example images
 examples = [
     ["example_image1.jfif"],
     ["example_image2.jfif"],
     ["example_image3.jfif"]
 ]
 
-# Set up the Gradio interface
+# Gradio interface
 iface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
-    outputs="text",  # Output will be text listing the predictions
-    title="Dental Image Classification",
-    description="Upload an image or select from the examples below to predict its class. The model accounts for the possibility of multiple dental flaws in a single image.",
-    examples=examples  # Add example images
+    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
 )
 
-# Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()