import gradio as gr
import yolov5
import os
from transformers import pipeline

imageClassifier = pipeline(task="image-classification",
                           model="Ara88/timri-model")

model = yolov5.load('./gentle-meadow.pt', device="cpu")


def predict(image):

    predictions = imageClassifier(image)
    maxScore = 0
    predictedLabel = None
    labelsForLocalization = ['meningioma', 'pituitary']
    output = {}
    for item in predictions:
        output[item['label']] = item['score']
        if (maxScore < item['score']):
            maxScore = item['score']
            predictedLabel = item['label']
    if (predictedLabel in labelsForLocalization):
        results = model([image], size=224)
        imageWithLocalization = results.render()[0]
    else:
        imageWithLocalization = image
    return output, imageWithLocalization


title = "Detecting Tumors in MRI Images"
description = """
  Try the examples at bottom to get started.
"""
examples = [
    [os.path.abspath('examples/sample_1.jpg')],
    [os.path.abspath('examples/sample_2.jpg')],
    [os.path.abspath('examples/sample_3.jpg')],
    [os.path.abspath('examples/sample_4.jpg')],
    [os.path.abspath('examples/sample_5.jpg')],
    [os.path.abspath('examples/sample_6.jpg')],
    [os.path.abspath('examples/sample_7.jpg')],
    [os.path.abspath('examples/sample_8.jpg')],
]

inputs = gr.Image(type="pil", shape=(224, 224),
                  label="Upload your image for detection")

outputs = [
    gr.Label(label="Tumor Classification"),
    gr.Image(type="pil", label="Tumor Detections")
]

interface = gr.Interface(
    fn=predict,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    description=description,
    cache_examples=True,
    theme='huggingface'
)
interface.launch(debug=True, enable_queue=True)