简介:
目标检测在生活中应用领域非常广泛,列如:道路违规抓拍、未戴口罩识别、工地未佩戴安全帽抓拍、厨房出现老鼠检测。 还可以用在游戏辅助外挂。以枪战为例,在游戏过程中时刻检测有没有人头出现。当检测到目标人头,再调用鼠标把枪口焦点移动到人头的位置,实现爆头效果。 本案例对yolov5官方代码进行了缩减,留下精华的部分,并封装成类。 yolov5默认模型支持80种目标检测,具体类型在文章最后。如果需要检测其他的目标,就需要自己收集数据进行数据标注,再重新训练新的模型。
1.参数初始化
def __init__(self):
# classes,模型训练的时候每种目标都有自己的一个标识,类型数字从0开始。这里不指定,默认全部检测。
self.classes = None
# 加载pt模型
self.weights = 'yolov5s.pt'
self.imgsz = [640, 640]
# 置信度,检测目标小于这个值的将不会被识别出来
self.conf_thres = 0.5
self.iou_thres = 0.45
# Load model
self.device = select_device('')
self.model = DetectMultiBackend(self.weights, device=self.device, dnn=False, fp16=False)
self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
self.imgsz = check_img_size(self.imgsz, s=self.stride)
self.model.warmup(imgsz=(1, 3, *self.imgsz))
2.算法推理
def detect(self, im0s):
img = letterbox(im0s, new_shape=(640, 640), stride=self.stride, auto=self.pt)[0]
img = img.transpose((2, 0, 1))[::-1]
img = np.ascontiguousarray(img)
im = torch.from_numpy(img).to(self.device)
im = im.half() if self.model.fp16 else im.float()
im /= 255
if len(im.shape) == 3:
im = im[None]
pred = self.model(im, augment=False, visualize=False)
pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, False, max_det=300)
dets = []
for i, det in enumerate(pred):
im0 = im0s.copy()
if len(det):
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
for *xyxy, conf, cls in reversed(det):
c = int(cls)
dets.append([int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]), self.names[c], f'{conf:.2f}'])
3.调用算法:
from detect import SmokeUtil
import cv2
su = SmokeUtil()
if __name__ == "__main__":
cap = cv2.VideoCapture('target.mp4')
while True:
success, frame = cap.read()
if not success:
break
dets = su.detect(frame)
if len(dets) > 0:
for j in dets:
cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (11, 44, 55), 3)
cv2.imshow('show', frame)
if cv2.waitKey(1) == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
frame = cv2.imread('1.jpg')
dets = su.detect(frame)
if len(dets) > 0:
for j in dets:
cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (0,255,0), 3)
cv2.imwrite('2.jpg', frame)
4.检测效果——以大象为检测目标:
(原图)
(检测效果)
5.视频流检测——车辆检测
6.支持检测类型
['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra','giraffe',
'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
'banana', 'apple', 'sandwich', 'orange', 'broccoli','carrot', 'hot dog', 'pizza',
'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
'hair drier', 'toothbrush']
|