import gradio as gr
from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
import torch
import os
from huggingface_hub import login

# Login to Hugging Face Hub
if 'HF_TOKEN' in os.environ:
    login(token=os.environ['HF_TOKEN'])

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"]
        ]
    )

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
demo.launch()