import gradio as gr
from transformers import pipeline

# مدل تشخیص بیماری پوستی از Hugging Face
clf = pipeline(
    task="image-classification",
    model="Anwarkh1/Skin_Cancer-Image_Classification"
)

# ترجمه لیبل‌ها
FA_LABELS = {
    "actinic_keratoses": "کراتوز آکتینیک (ضایعه پیش‌سرطانی)",
    "basal_cell_carcinoma": "سرطان سلول بازال",
    "benign_keratosis-like_lesions": "کراتوز خوش‌خیم",
    "dermatofibroma": "درماتوفیبروما (توده پوستی خوش‌خیم)",
    "melanoma": "ملانوما (سرطان پوست)",
    "nv": "خال معمولی",
    "vascular_lesions": "ضایعات عروقی",
}

def normalize(label: str):
    return (
        label.strip().lower()
        .replace(" ", "_")
    )

def translate_label(label: str):
    key = normalize(label)
    return FA_LABELS.get(key, label)

def predict(image):
    preds = clf(image)
    out = {}
    for p in preds:
        fa_label = translate_label(p["label"])
        out[fa_label] = float(p["score"])
    return out

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="تصویر پوست را آپلود کنید"),
    outputs=gr.Label(num_top_classes=3, label="نتیجه مدل"),
    title="تشخیص ضایعات پوستی (نسخه آموزشی)",
    description="این ابزار فقط برای آموزش ساخته شده و تشخیص پزشکی قطعی نیست."
)

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