import gradio as gr
import numpy as np
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

# Load model + processor
processor = AutoImageProcessor.from_pretrained("prithivMLmods/Weather-Image-Classification")
model = AutoModelForImageClassification.from_pretrained("prithivMLmods/Weather-Image-Classification")

def classify_weather(image_input):
    try:
        if isinstance(image_input, np.ndarray):
            image = Image.fromarray(image_input.astype("uint8")).convert("RGB")
        else:
            raise TypeError("Only NumPy array input is supported.")

        # preprocess as batch
        inputs = processor(images=[image], return_tensors="pt")

        # inference
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class_id = logits.argmax(-1).item()
            predicted_label = model.config.id2label[predicted_class_id]

            # optional: return probabilities for Label(num_top_classes=5)
            probs = torch.softmax(logits, dim=-1).squeeze().tolist()
            labels = [model.config.id2label[i] for i in range(len(probs))]
            output_dict = dict(zip(labels, probs))

        return output_dict

    except Exception as e:
        return {"Error": str(e)}


# Gradio interface
iface = gr.Interface(
    fn=classify_weather,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(num_top_classes=5, label="Weather Condition"),
    title="Weather Image Classification",
    description="Upload an image to classify the weather condition (sun, rain, snow, fog, or clouds)."
)

if __name__ == "__main__":
    iface.launch(show_error=True)