File size: 2,859 Bytes
e1e6433 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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']
# for i, r in enumerate(results):
# # Plot results image
# im_bgr = r.plot()
# im_rgb = im_bgr[..., ::-1] # Convert BGR to RGB
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'
# im_rgb = Image.fromarray(im_rgb)
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)
# xywh = int(results.boxes.xywh)
# x = xywh[0]
# y = xywh[1]
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")
]
# 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]
# ]
iface = gr.Interface(fn=response2, inputs=inputs, outputs=outputs)
iface.launch(debug=True) |