# coding: utf8
import gradio as gr
from transformers import pipeline

pipeline = pipeline(task="image-classification", model="dqnguyen/Diabetic_Foot_Ulcer_Image_Classification")

def predict(image):
    predictions = pipeline(image)
    #return {p["label"]: p["score"] for p in predictions}
    results = {}
    for p in predictions:
        if p["label"] == "MoHat":
            results["Granulation tissue (Mô hạt)"] = p["score"]
        elif p["label"] == "MoGiaMacNhiemKhuan":
            results["Pseudomembranous tissue with a bacterial infection (Mô giả mạc nhiễm khuẩn)"] = p["score"]
        elif p["label"] == "MoHoaiTu":
            results["Necrotic tissue (Mô hoại tử)"] = p["score"]
    return results
    
gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Upload an image (Tải một bức ảnh vết loét tiểu đường)", type="filepath"),
    outputs=gr.outputs.Label(num_top_classes=5),
    title="Diabetic Ulcer Image Classification (Phân loại ảnh vết loét tiểu đường)",
).launch()