import gradio as gr
from transformers import pipeline
import torch
from torchvision import transforms as T

def calc_result_confidence (model_output):
    probs = torch.nn.functional.softmax(model_output, dim=1)
    conf, classes = torch.max(probs, 1)

    return conf.item(), classes.item()

def downsyndrome_gradio_inference(img_file):
    classes = ['Down Syndrome', 'Healty']
    infer_transform = T.Compose([
        T.Resize((255, 255)),
        T.ToTensor(),
    ])
    transform_image = infer_transform(img_file.convert('RGB')).float().unsqueeze(0)
    model = pipeline(task='image-classification', model='gitfreder/down-syndrome-detection')
    conf, cls = calc_result_confidence(model(transform_image))
    
    return {
        'Predicted': classes[cls], 
        'Confidence Score': conf
        }


iface = gr.Interface(fn=downsyndrome_gradio_inference, inputs=gr.Image(type='pil'), outputs=gr.JSON(), title="Down Syndrome Detection", description="A model interfaces that detect downsyndrom children from the photo")
iface.launch()