|
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) |
|
|