import gradio as gr
import os
import torch


from model import predict


class_names = [ 'benign_keratosis-like_lesions','basal_cell_carcinoma','actinic_keratoses','dermatofibroma','melanocytic_Nevi']
example_names = ["actinic_keratoses","basal_cell_carcinoma","melanocytic_Nevi"]

title = "Skin Cancer Classifier"
description = "An ViT computer vision model to classify images from HAM10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. <br/> List: benign_keratosis-like_lesions, basal_cell_carcinoma, actinic_keratoses, dermatofibroma, melanocytic_Nevi"
# article = "https://dataverse.harvard.edu/dataset.xhtml?persistentId=doi:10.7910/DVN/DBW86T"

# Create examples list from "examples/" directory
example_list = [["examples/" + example, example.split('_')[0]] for example in os.listdir("examples")]
# print(example_list)
# result , timing = predict(example_list[0])
# 

# Create a single dictionary


# Output the combined dictionary
# print(combined_dict)

# Create the Gradio demo
# The output of the prediction must be in a dictionary format!
demo = gr.Interface(
    fn=predict,  # mapping function from input to output
    inputs=gr.Image(type="pil"),  # what are the inputs?
    outputs=[gr.Label(num_top_classes=5, label="Predictions"),  # what are the outputs?
             gr.Number(label="Prediction time (s)")], 
    examples=example_list, 
    title=title,
    description=description,
    example_labels=example_names
)

# Launch the demo!
demo.launch()