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

imageClassifier = pipeline(task="image-classification",
                           model="PranomVignesh/Police-vs-Public")


def predict(image):
    model = yolov5.load('./best.pt', device="cpu")
    results = model([image], size=224)

    predictions = imageClassifier(image)
    classMappings = {
        'police': "Police / Authorized Personnel",
        'public': 'Unauthorized Person'
    }
    output = {}
    for item in predictions:
        output[classMappings[item['label']]] = item['score']

    return results.render()[0], output


title = "Detecting Unauthorized Individuals with Firearms"

examples = [
    []
]

title = "Detecting Unauthorized Individuals with Firearms"
description = """
  Try the examples at bottom to get started.
"""
examples = [[
    os.path.join(os.path.abspath(''), './examples/sample_1.png'),
    os.path.join(os.path.abspath(''), './examples/sample_2.png'),
    os.path.join(os.path.abspath(''), './examples/sample_3.png'),
    os.path.join(os.path.abspath(''), './examples/sample_4.png'),
    os.path.join(os.path.abspath(''), './examples/sample_5.png'),
    os.path.join(os.path.abspath(''), './examples/sample_6.png'),
    os.path.join(os.path.abspath(''), './examples/sample_7.png'),
    os.path.join(os.path.abspath(''), './examples/sample_8.png'),
]]

inputs = gr.Image(type="pil", shape=(224, 224),
                  label="Upload your image for detection")
outputs = [
    gr.Image(type="pil", label="Gun Detections"),
    gr.Label(label="Class Prediction")
]

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