写一段python基于yolov5人脸识别的代码
时间: 2023-10-29 20:06:57 浏览: 53
好的,以下是一个基于yolov5的人脸识别代码示例:
```python
import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
def detect_faces(img_path, conf_thresh=0.5, iou_thresh=0.45, device='cpu'):
# 加载模型
model = attempt_load('yolov5s.pt', map_location=device)
img_size = model.img_size
stride = int(model.stride.max()) # 均匀采样间隔
# 图像预处理
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w = img.shape[:2]
img_tensor = torch.from_numpy(img).to(device).unsqueeze(0)
img_tensor = img_tensor.float() / 255.0
img_tensor = img_tensor.permute(0, 3, 1, 2)
# 模型推理
pred = model(img_tensor)[0]
pred = non_max_suppression(pred, conf_thresh, iou_thresh, classes=[0])
# 处理预测结果
bboxes = []
confs = []
for det in pred:
if det is not None and len(det):
det[:, :4] = scale_coords(img_tensor.shape[2:], det[:, :4], (h, w)).round()
for *xyxy, conf, cls in reversed(det):
if cls == 0:
bboxes.append(xyxy)
confs.append(conf)
return bboxes, confs
if __name__ == '__main__':
# 检测人脸
img_path = 'test.jpg'
device = select_device('0')
bboxes, confs = detect_faces(img_path, device=device)
# 显示结果
img = cv2.imread(img_path)
for bbox, conf in zip(bboxes, confs):
x1, y1, x2, y2 = bbox
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.putText(img, f'{conf:.2f}', (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
cv2.imshow('result', img)
cv2.waitKey()
```
这个代码可以检测一张图片中的人脸,并将检测结果以矩形框的形式显示在图片上。需要注意的是,这个代码只能检测单张图片,如果需要批量处理多张图片,还需要进一步修改。