from transformers import ViTFeatureExtractor, ViTForImageClassification
from PIL import Image
import torch
import gradio as gr

# Model en extractor laden
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")

# Classificatiefunctie
def classify_space_image(image):
    inputs = extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted = logits.argmax(-1).item()
    label = model.config.id2label[predicted]
    return f"Voorspelling: {label}"

# Gradio Interface
demo = gr.Interface(
    fn=classify_space_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="🔭 AstroClassifier",
    description="Upload een ruimte-afbeelding (bijv. van NASA of een telescoop) en ontdek of het een ster, planeet, nevel of zwart gat is (volgens de AI!)."
)

demo.launch()