import gradio as gr
from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch

model_name = "Lalith47/custom-cloud-model"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

def classify_image(image):
    """Returns ALL predictions sorted by confidence"""
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
    
    # Get all predictions with their labels
    all_predictions = []
    for idx, prob in enumerate(probs):
        label = model.config.id2label[idx]
        all_predictions.append([label, float(prob)])
    
    # Sort by confidence (highest first)
    all_predictions.sort(key=lambda x: x[1], reverse=True)
    
    # Print top 2
    print(f"🥇 Top: {all_predictions[0][0]} ({all_predictions[0][1]*100:.1f}%)")
    print(f"🥈 2nd: {all_predictions[1][0]} ({all_predictions[1][1]*100:.1f}%)")
    
    return all_predictions

iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="☁️ Cloud Classifier",
    description="Upload a cloud image to identify its type.",
    live=False
)

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