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

# Hugging Face tokeninizi çevresel değişkenden alın
hf_token = os.getenv("HF_TOKEN")
if not hf_token:
    raise ValueError("HF_TOKEN çevresel değişkeni ayarlanmamış. Lütfen Hugging Face token'ınızı ayarlayın.")

# Model ve işlemciyi yükleyin
model_name = "cantuncok/autotrain1-model"
processor = AutoImageProcessor.from_pretrained(model_name, token=hf_token)
model = AutoModelForImageClassification.from_pretrained(model_name, token=hf_token)

# Görsel sınıflandırma fonksiyonu
def classify_image(img):
    if img is None:
        return {"Hata": 1.0}
    
    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    
    labels = model.config.id2label
    predicted_label = labels[predicted_class_idx]
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    return {label: float(probabilities[0, idx]) for idx, label in labels.items()}

# Gradio arayüzü
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Görüntü Sınıflandırma Uygulaması",
    description="Bu uygulama, 'cantuncok/autotrain1-model' modeli kullanılarak görüntü sınıflandırması yapar."
)

# Uygulamayı başlatma
if __name__ == "__main__":
    iface.launch()