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

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


model = torch.hub.load(
    'ultralytics/yolov5',
    'custom',
    path='./best.pt',
    device="cpu",
    force_reload=True
)
model.eval()


def predict(image):
    results = model([image], size=224)
    print(results)
    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"
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.jpg')],
    [os.path.join(os.path.abspath(''), './examples/sample_4.jpg')],
    [os.path.join(os.path.abspath(''), './examples/sample_5.jpg')],
    [os.path.join(os.path.abspath(''), './examples/sample_6.jpg')],
    [os.path.join(os.path.abspath(''), './examples/sample_7.jpg')],
    [os.path.join(os.path.abspath(''), './examples/sample_8.jpg')],
]

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,
    theme='huggingface'
)
interface.launch(debug=True, enable_queue=True)