如何把yolov5检测不出来的照片单独存放在一个文件夹里
时间: 2024-02-20 16:59:38 浏览: 173
您可以使用Python和YoloV5的API来实现将YoloV5检测不出来的照片单独存放在一个文件夹中的功能。以下是实现的步骤:
1. 导入必要的库和模块:
```
import os
import shutil
import torch
import cv2
from models.experimental import attempt_load
from utils.datasets import LoadImages
from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized
```
2. 加载YoloV5模型:
```
weights = 'yolov5s.pt'
device = select_device('')
model = attempt_load(weights, map_location=device)
imgsz = check_img_size(640, s=model.stride.max())
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
```
3. 定义检测函数:
```
def detect(img, conf_thres=0.25, iou_thres=0.45):
img = torch.from_numpy(img).to(device)
img = img.float()
img /= 255.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
pred = model(img)[0]
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=None, agnostic=False)
if pred[0] is None:
return None
boxes = pred[0][:, :4].cpu().numpy()
scores = pred[0][:, 4].cpu().numpy()
classes = pred[0][:, 5].cpu().numpy().astype(int)
return boxes, scores, classes
```
4. 加载图片并进行检测:
```
image_folder = 'path/to/image/folder'
result_folder = 'path/to/result/folder'
os.makedirs(result_folder, exist_ok=True)
img_files = LoadImages(os.path.join(image_folder, '*'))
for img_i, (path, img, im0s, vid_cap) in enumerate(img_files):
boxes, scores, classes = detect(img)
if boxes is None:
shutil.move(path, os.path.join(result_folder, os.path.basename(path)))
continue
```
5. 将检测不出来的图片移动到指定文件夹:
```
for box, score, cls in zip(boxes, scores, classes):
if score < conf_thres:
shutil.move(path, os.path.join(result_folder, os.path.basename(path)))
break
```
6. 运行完整代码并将结果保存到指定文件夹:
```
python detect.py
```
其中,`detect.py`是包含完整代码的Python文件。
阅读全文