Gun_detection / app.py
NKASG's picture
Update app.py
81d9d84 verified
raw
history blame
1.39 kB
from ultralytics import YOLO
import cv2
import cvzone
import math
import imutils
from imutils.video import VideoStream
cap = cv2.VideoCapture(0)#'rtsp://admin:[email protected]:554/Streaming/Channels/1')
cap.set(3, 640)
cap.set(4, 128)
#cap = cv2.VideoCapture("C:/Users/Moh/PycharmProjects/pythonProject/video/GVD.mp4")
model = YOLO("best20.pt")
classNames = ['person', 'Gun' ]
while True:
success, frame = cap.read()
result = model(frame, stream=True)
for r in result:
boxes = r.boxes
for box in boxes:
# bounding box
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
w, h = x2 - x1, y2 - y1
#confidence
conf = math.ceil((box.conf[0] * 100)) / 100
#className
cls = int(box.cls[0])
currentClass = classNames[cls]
if currentClass == "Gun" and conf > 0.4:
print(conf)
cvzone.cornerRect(frame, (x1, y1, w, h), l=9, rt=5, colorR=(255, 0, 0))
cvzone.putTextRect(frame, f'{conf}{currentClass}', (max(0, x1), max(0, y1)))
print(conf)
frame = imutils.resize(frame, width=1200)
cv2.imshow('AsimCodeCam', frame)
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
break
cv2.destroyAllWindows()
video_stream.stop()