import gradio as gr
from image_classifier import ImageClassifier
import numpy as np

classifier = ImageClassifier()

def classify_image(image):
    # Convert Gradio image to the format expected by our classifier
    # Our classifier expects a file path or URL, so we'll save the image temporarily
    import tempfile
    import os
    from PIL import Image
    
    # Save the image temporarily
    with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
        if isinstance(image, np.ndarray):
            # Convert numpy array to PIL Image
            pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
        else:
            pil_img = image
        pil_img.save(tmp.name)
        tmp_path = tmp.name
    
    # Classify the image
    try:
        results = classifier.classify_image(tmp_path)
        # Format results for display
        labels = [res['label'] for res in results]
        confidences = [res['probability'] for res in results]
        
        # Clean up temporary file
        os.remove(tmp_path)
        
        return labels, confidences
    except Exception as e:
        # Clean up temporary file even if there's an error
        os.remove(tmp_path)
        raise e

demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload an image for classification"),
    outputs=[
        gr.Label(num_top_classes=5, label="Top Predictions"),
        gr.BarPlot(x="Label", y="Confidence", title="Confidence Scores", width=500, height=300)
    ],
    title="🖼️ Computer Vision Model",
    description="This model performs image classification using a pre-trained ResNet model.",
    examples=[]
)

if __name__ == "__main__":
    demo.launch()