yolov7将视频中检测含A类且不含B、C、D类的截图保存到A文件夹,将同时检测到B、C类的截图到B文件夹,将检测到D的保存到D文件夹
时间: 2023-05-12 09:04:51 浏览: 72
可以使用以下代码实现:
```python
import cv2
import numpy as np
# Load YOLOv7 model
net = cv2.dnn.readNet("yolov7.weights", "yolov7.cfg")
# Define classes
classes = ["A", "B", "C", "D"]
# Define colors for each class
colors = np.random.uniform(0, 255, size=(len(classes), 3))
# Load video
cap = cv2.VideoCapture("video.mp4")
# Define output folders
folder_a = "A"
folder_b = "B"
folder_d = "D"
# Loop through frames
while True:
ret, frame = cap.read()
if not ret:
break
# Detect objects in frame
blob = cv2.dnn.blobFromImage(frame, 1/255, (416, 416), swapRB=True)
net.setInput(blob)
outs = net.forward(net.getUnconnectedOutLayersNames())
# Extract class, confidence and bounding box for each detection
class_ids = []
confidences = []
boxes = []
for out in outs:
for detection in out:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.5:
center_x = int(detection[0] * frame.shape[1])
center_y = int(detection[1] * frame.shape[0])
width = int(detection[2] * frame.shape[1])
height = int(detection[3] * frame.shape[0])
left = int(center_x - width/2)
top = int(center_y - height/2)
class_ids.append(class_id)
confidences.append(float(confidence))
boxes.append([left, top, width, height])
# Apply non-maximum suppression to remove overlapping detections
indices = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
# Save images to corresponding folders
for i in indices:
i = i[0]
class_id = class_ids[i]
if classes[class_id] == "A" and "B" not in classes and "C" not in classes and "D" not in classes:
cv2.imwrite(f"{folder_a}/frame_{i}.jpg", frame)
elif classes[class_id] == "B" and "C" in classes:
cv2.imwrite(f"{folder_b}/frame_{i}.jpg", frame)
elif classes[class_id] == "D":
cv2.imwrite(f"{folder_d}/frame_{i}.jpg", frame)
# Release video and close windows
cap.release()
cv2.destroyAllWindows()
```
其中,`yolov7.weights` 和 `yolov7.cfg` 是预训练的 YOLOv7 模型文件,需要先下载并保存到本地。`classes` 定义了模型需要检测的类别,`colors` 定义了每个类别对应的颜色。`cap` 是视频对象,`folder_a`、`folder_b` 和 `folder_d` 分别是保存截图的文件夹。在每一帧中,我们首先使用 YOLOv7 模型检测出所有的目标,并根据类别将截图保存到对应的文件夹中。
相关推荐
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![zip](https://img-home.csdnimg.cn/images/20210720083736.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)
![](https://csdnimg.cn/download_wenku/file_type_ask_c1.png)