python简单用opencv打开摄像头并用yolov5模型进行物体检测
·
python简单用opencv打开摄像头并用yolov5模型进行物体检测
1.GitHub代码
yolov5:https://github.com/ultralytics/yolov5
2.环境准备
pip install -r requirements.txt
3.示例代码
import torch
# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # or yolov5m, yolov5l, yolov5x, custom
# Images
img = 'https://ultralytics.com/images/zidane.jpg' # or file, Path, PIL, OpenCV, numpy, list
# Inference
results = model(img)
# Results
results.print() # or .show(), .save(), .crop(), .pandas(), etc.
或者
使用命令行
python detect.py --source 0 # webcam
#file.jpg # image
#file.mp4 # video
#path/ # directory
#path/*.jpg # glob
#'https://youtu.be/NUsoVlDFqZg' # YouTube
#'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
4.个人修改简单代码(直接通过opencv打开摄像头进行识别)
# author:chouti
import sys
from pathlib import Path
import cv2
import torch
FILE = Path(__file__).absolute()
sys.path.append(FILE.parents[0].as_posix()) # add yolov5/ to path
from models.experimental import attempt_load
from utils.general import is_ascii, non_max_suppression, scale_coords, set_logging
from utils.plots import Annotator, colors
from utils.torch_utils import select_device
@torch.no_grad()
def run(weights='yolov5s.pt', # model.pt path(s)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
line_thickness=3, # bounding box thickness (pixels)
half=False, # use FP16 half-precision inference
):
# Initialize
set_logging()
device = select_device(device)
print(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
model = attempt_load(weights, map_location=device) # load FP32 model
names = model.module.names if hasattr(model, 'module') else model.names # get class names
ascii = is_ascii(names) # names are ascii (use PIL for UTF-8)
capture = cv2.VideoCapture(0)
while True:
# 获取一帧
ret, frame = capture.read()
img = torch.from_numpy(frame).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img = img / 255.0 # 0 - 255 to 0.0 - 1.0
if len(img.shape) == 3:
img = img[None] # expand for batch dim
img = img.transpose(2, 3)
img = img.transpose(1, 2)
# Inference
pred = model(img, augment=False, visualize=False)[0]
# NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
# Process predictions
for i, det in enumerate(pred): # detections per image
s = ''
annotator = Annotator(frame, line_width=line_thickness, pil=not ascii)
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += str(n.item()) + ' ' + str(names[int(c)]) + ' ' # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
c = int(cls) # integer class
label = f'{names[c]} {conf:.2f}'
annotator.box_label(xyxy, label, color=colors(c, True))
print(xyxy)
print('result:'+s)
cv2.imshow('frame', frame)
if cv2.waitKey(1) == ord('q'):
break
def main():
run()
if __name__ == "__main__":
main()
5.注意事项
①可能需要提前下载好模型
②另外可能需要GPU
中国需要更多的免费开源!
更多推荐
已为社区贡献1条内容
所有评论(0)