import gradio as gr
import os
from utils import page_utils

from ultralytics import YOLO

# Load a model
model = YOLO('model_- 14 december 2023 12_01.pt')  # pretrained YOLOv8n model

class_names = ['abdominal', 'adult', 'others', 'pediatric', 'spine']
class_names.sort()

examples_dir = "samples"

def image_classifier(inp):
    """Image Classifier Function.
    Parameters
    ----------
    inp: Optional[np.ndarray] = None
        Input image from callback
    Returns
    -------
    Dict
        A dictionary class names and its probability
    """

    # If input not valid, return dummy data or raise error
    if inp is None:
        return {'cat': 0.3, 'dog': 0.7}
    
    result = model(inp)

    # postprocess
    labeled_result = {class_names[label]: confidence for label, confidence in zip(result.probs.top5, result.probs.top5conf)}

    return labeled_result

# gradio code block for input and output
with gr.Blocks() as app:
    gr.Markdown("# Lung Cancer Classification")

with open('index.html', encoding="utf-8") as f:
    description = f.read()

# gradio code block for input and output
with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR, secondary_hue=page_utils.KALBE_THEME_COLOR).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_500",
        button_primary_text_color="white",
    )) as app:
    with gr.Column():
        gr.HTML(description)

    with gr.Row():
        with gr.Column():
            inp_img = gr.Image()
            with gr.Row():
                clear_btn = gr.Button(value="Clear")
                process_btn = gr.Button(value="Process", variant="primary")
        with gr.Column():
            out_txt = gr.Label(label="Probabilities", num_top_classes=5)

    process_btn.click(image_classifier, inputs=inp_img, outputs=out_txt)
    clear_btn.click(lambda:(
        gr.update(value=None),
        gr.update(value=None)
        ),
        inputs=None,
        outputs=[inp_img, out_txt])

    gr.Markdown("## Image Examples")
    gr.Examples(
        examples=[os.path.join(examples_dir, "1.2.840.113564.1921681202.202011100756242032.1203801020003.dcm.jpeg")
                 ],
        inputs=inp_img,
        outputs=out_txt,
        fn=image_classifier,
        cache_examples=False,
    )
    gr.Markdown(line_breaks=True, value='Author: Jason Adrian (jasonadriann6@gmail.com) <div class="row"><a href="https://github.com/jasonadriann?tab=repositories"><img alt="GitHub" src="https://img.shields.io/badge/Jason%20Adrian-000000?logo=github"> </div>')

# demo = gr.Interface(fn=image_classifier, inputs="image", outputs="label")
app.launch(share=True)