Update app.py

app.py

@@ -1,45 +1,75 @@
 import gradio as gr
 from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
 import torch
+import os
+from huggingface_hub import login
 
-# Initialize the model and processor
-processor = AutoImageProcessor.from_pretrained("alexdekan030/autotrain-1c9rx-gdx3r")
-model = AutoModelForImageClassification.from_pretrained("alexdekan030/autotrain-1c9rx-gdx3r")
-pipe = pipeline("image-classification", model=model, image_processor=processor)
+# Login to Hugging Face Hub
+if 'HF_TOKEN' in os.environ:
+    login(token=os.environ['HF_TOKEN'])
 
-def predict_pneumonia(image):
-    """
-    Predict whether an image shows pneumonia or normal chest X-ray
-    Args:
-        image: Input image
-    Returns:
-        dict: Dictionary containing prediction probabilities
-    """
-    # Make prediction
-    result = pipe(image)
-    
-    # Create a formatted output dictionary
-    probabilities = {pred['label']: float(pred['score']) for pred in result}
-    
-    return probabilities
+try:
+    # Initialize the model and processor
+    processor = AutoImageProcessor.from_pretrained(
+        "alexdekan030/autotrain-awcru-nr8j7",
+        use_auth_token=os.environ.get('HF_TOKEN')
+    )
+    model = AutoModelForImageClassification.from_pretrained(
+        "alexdekan030/autotrain-awcru-nr8j7",
+        use_auth_token=os.environ.get('HF_TOKEN')
+    )
+    pipe = pipeline("image-classification", model=model, image_processor=processor)
+
+    def predict_pneumonia(image):
+        """
+        Predict whether an image shows pneumonia or normal chest X-ray
+        Args:
+            image: Input image
+        Returns:
+            dict: Dictionary containing prediction probabilities
+        """
+        if image is None:
+            return {"Error": 1.0}
+            
+        try:
+            # Make prediction
+            result = pipe(image)
+            
+            # Create a formatted output dictionary
+            probabilities = {pred['label']: float(pred['score']) for pred in result}
+            
+            return probabilities
+        except Exception as e:
+            return {"Error": 1.0}
 
-# Create the Gradio interface
-demo = gr.Interface(
-    fn=predict_pneumonia,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=2),
-    title="Pneumonia Detection from Chest X-rays",
-    description="""Upload a chest X-ray image to detect if it shows signs of pneumonia.
-                   The model will classify the image as either 'NORMAL' or 'PNEUMONIA'
-                   and provide confidence scores for each class.""",
-    examples=[
-        # You can add example images here if you have them
-        # ["path/to/example1.jpg"],
-        # ["path/to/example2.jpg"]
-    ],
-    theme=gr.themes.Base()
-)
+    # Create the Gradio interface
+    demo = gr.Interface(
+        fn=predict_pneumonia,
+        inputs=gr.Image(type="pil"),
+        outputs=gr.Label(num_top_classes=2),
+        title="Pneumonia Detection from Chest X-rays",
+        description="""Upload a chest X-ray image to detect if it shows signs of pneumonia.
+                       The model will classify the image as either 'NORMAL' or 'PNEUMONIA'
+                       and provide confidence scores for each class.""",
+        examples=[
+            # You can add example images here if you have them
+            # ["path/to/example1.jpg"],
+            # ["path/to/example2.jpg"]
+        ]
+    )
+
+except Exception as e:
+    # Create a simple interface if model loading fails
+    def error_interface(image):
+        return {"Error": "Model failed to load. Please check authentication and model availability."}
+    
+    demo = gr.Interface(
+        fn=error_interface,
+        inputs=gr.Image(type="pil"),
+        outputs=gr.Label(num_top_classes=1),
+        title="Error Loading Model",
+        description="There was an error loading the model. Please check if the model is accessible and authentication is correct."
+    )
 
 # Launch the app
-if __name__ == "__main__":
-    demo.launch(share=True)
+demo.launch()