Spaces:
Sleeping
Sleeping
File size: 1,388 Bytes
c640bc9 066c36c c640bc9 066c36c c640bc9 066c36c c640bc9 066c36c c640bc9 |
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 |
"""
Created By: ishwor subedi
Date: 2024-04-04
"""
import cv2 as cv
from src.services.weapon_det_service.weapon_detection_service import DetectionService
def single_image_inference(image_path):
"""
This function is used to detect the weapon in the image and save the image with bounding box
:param image_path: image path to detect the weapon
:return: image vector with bounding box
"""
detection_service = DetectionService(
model_path='resources/models/v2/best.pt',
)
results = detection_service.image_det_save(
image_path=image_path,
thresh=0.4
)
for result in results:
original_image = result.orig_img
bbox = result.boxes.xyxy.int().tolist()
for i, bbox in enumerate(bbox):
x1, y1, x2, y2 = bbox
cv.putText(original_image, f'{result.names[result.boxes.cls[i].int().tolist()]}', (x1, y1 - 10),
cv.FONT_HERSHEY_SIMPLEX,
0.9,
(0, 0, 255), 2)
cv.rectangle(original_image, (x1, y1), (x2, y2), (0, 0, 255), 2)
return original_image
if __name__ == '__main__':
image_path = '/home/ishwor/Desktop/gun-detection/images/cam_images/th-3711382641.jpg'
image = single_image_inference(image_path)
cv.imshow('Weapon Detection', image)
cv.waitKey(0)
cv.destroyAllWindows()
|