Update app.py

app.py

@@ -1,30 +1,51 @@
-from huggingface_hub import hf_hub_download
+import os
 import torch
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
 from transformers import ViTImageProcessor
 from PIL import Image
 import gradio as gr
-from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from huggingface_hub import hf_hub_download
 
 # Repository configuration
 REPO_ID = "IFMedTech/Dental_Q"
 MODEL_FILENAME = "quantized_model.ptl"
 
-# Download the model file from Hugging Face Hub
-model_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename=MODEL_FILENAME
-)
+def download_model_from_hub():
+    """Download model from private Hugging Face repository"""
+    token = os.environ.get("HUGGINGFACE_TOKEN")
+    
+    if not token:
+        raise ValueError(
+            "HUGGINGFACE_TOKEN environment variable is required for private repo access. "
+            "Please set it in your Space settings under 'Repository secrets'."
+        )
+    
+    try:
+        model_path = hf_hub_download(
+            repo_id=REPO_ID,
+            filename=MODEL_FILENAME,
+            token=token
+        )
+        return model_path
+    except Exception as e:
+        raise RuntimeError(f"Failed to download model from {REPO_ID}: {str(e)}")
 
-# Download the processor files (assuming they're in the same repo)
-processor_path = hf_hub_download(
-    repo_id=REPO_ID,
-    filename="preprocessor_config.json"
-)
+def load_model_and_processor():
+    """Load the model and processor"""
+    token = os.environ.get("HUGGINGFACE_TOKEN")
+    
+    # Download and load model
+    model_path = download_model_from_hub()
+    quantized_model = torch.jit.load(model_path, map_location="cpu")
+    quantized_model.eval()
+    
+    # Load processor from private repo
+    processor = ViTImageProcessor.from_pretrained(REPO_ID, token=token)
+    
+    return quantized_model, processor
 
-# Load Processor & Quantized Model
-processor = ViTImageProcessor.from_pretrained(REPO_ID)
-quantized_model = torch.jit.load(model_path, map_location="cpu")
-quantized_model.eval()
+# Initialize model and processor
+quantized_model, processor = load_model_and_processor()
 
 # Define Inference Preprocessing
 size = processor.size['height']
@@ -43,60 +64,4 @@ except AttributeError:
     label_names = ["Background", "Caries", "Normal Teeth", "Plaque"]
 
 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
-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()
+    """Load and preprocess a