|
import gradio as gr |
|
import matplotlib.pyplot as plt |
|
from PIL import Image |
|
from ultralyticsplus import YOLO, render_result |
|
import cv2 |
|
import numpy as np |
|
from transformers import pipeline |
|
|
|
model = YOLO('best (1).pt') |
|
model2 = pipeline('image-classification','Kaludi/csgo-weapon-classification') |
|
name = ['grenade','knife','pistol','rifle'] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def response(image): |
|
print(image) |
|
results = model(image) |
|
text = "" |
|
name_weap = "" |
|
|
|
for r in results: |
|
conf = np.array(r.boxes.conf) |
|
cls = np.array(r.boxes.cls) |
|
cls = cls.astype(int) |
|
xywh = np.array(r.boxes.xywh) |
|
xywh = xywh.astype(int) |
|
|
|
for con, cl, xy in zip(conf, cls, xywh): |
|
cone = con.astype(float) |
|
conef = round(cone,3) |
|
conef = conef * 100 |
|
text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n") |
|
|
|
if cl == 0: |
|
name_weap += name[cl] + '\n' |
|
elif cl == 1: |
|
name_weap += name[cl] + '\n' |
|
elif cl == 2: |
|
out = model2(image) |
|
name_weap += out[0]["label"] + '\n' |
|
elif cl == 3: |
|
out = model2(image) |
|
name_weap += out[0]["label"] + '\n' |
|
|
|
|
|
|
|
|
|
|
|
return name_weap, text |
|
|
|
|
|
|
|
def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6): |
|
|
|
results = model.predict(image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size) |
|
|
|
box = results[0].boxes |
|
|
|
render = render_result(model=model, image=image, result=results[0], rect_th = 1, text_th = 1) |
|
|
|
|
|
weapon_name, text_detection = response(image) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return render, text_detection, weapon_name |
|
|
|
|
|
inputs = [ |
|
gr.Image(type="filepath", label="Input Image"), |
|
gr.Slider(minimum=320, maximum=1280, value=640, |
|
step=32, label="Image Size"), |
|
gr.Slider(minimum=0.0, maximum=1.0, value=0.3, |
|
step=0.05, label="Confidence Threshold"), |
|
gr.Slider(minimum=0.0, maximum=1.0, value=0.6, |
|
step=0.05, label="IOU Threshold"), |
|
] |
|
|
|
|
|
outputs = [gr.Image( type="filepath", label="Output Image"), |
|
gr.Textbox(label="Result"), |
|
gr.Textbox(label="Weapon Name") |
|
] |
|
|
|
title = "YOLOv8 Custom Object Detection by Uyen Nguyen" |
|
|
|
|
|
examples = [['th (11).jpg', 640, 0.3, 0.6], |
|
['th (8).jpg', 640, 0.3, 0.6], |
|
['th (3).jpg', 640, 0.3, 0.6], |
|
['th.jpg', 640, 0.15, 0.6], |
|
['th (2).jpg', 640, 0.15, 0.6], |
|
] |
|
|
|
|
|
|
|
|
|
iface = gr.Interface(fn=response2, inputs=inputs, outputs=outputs, examples=examples) |
|
iface.launch(debug=True) |
|
|
