File size: 2,971 Bytes
aeeb601 afc670e b6d54f1 b0e1d53 662f035 b0e1d53 4862689 7a8c28d 662f035 4862689 fd0e675 4862689 6adade1 662f035 8f46261 6d29c3c 66d9a45 aa4c587 1425272 662f035 6adade1 662f035 6adade1 662f035 39df642 239fc8b 39df642 662f035 4862689 be96225 c134122 6adade1 662f035 c134122 81b89b8 53a60d6 e232795 939d079 b6d54f1 3f46ab4 78dbbdb bb013d8 662f035 4862689 07456a8 b7b81f5 09410a4 4862689 78dbbdb 662f035 b6d54f1 c1c1237 b6d54f1 6e42dcd b6d54f1 6e42dcd b6d54f1 7ba0645 662f035 7ba0645 b6d54f1 1e51cca b1ec0c4 94ef2cb 662f035 b1ec0c4 b6d54f1 c134122 33110c9 24d7427 |
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 106 107 108 109 |
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")
]
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)
|