当前位置：首页 > news >正文

视频分帧【截取图片】（YOLO目标检测【生成数据集】）

news 2025/8/23 8:28:28

高效率制作数据集【按这个流程走，速度很顶】

本次制作，1059张图片【马路上流动车辆】
几乎就是全自动了，只要视频拍得好，YOLO辅助制作数据集就效率极高

视频中的图片抽取：
【由于视频内存过大，遇到报错执行失败，解决方法已附加在代码下面】

import cv2
import os
import pdb
import numpy as np
#from glob2 import globvideos_src_path = 'F:\\testkk\\vivivi\\'  # 提取图片的视频文件夹# 筛选文件夹下MP4格式的文件
# videos = os.listdir(videos_src_path)  # 用于返回指定的文件夹包含的文件或文件夹的名字的列表。
# videos = filter(lambda x: x.endswith('mp4'), videos)
dirs = os.listdir(videos_src_path)  # 获取指定路径下的文件
count = 0
# 写入txt
f = "F:\\testkk\\images\\data.txt"
with open(f, "w+") as file:file.write("-----start-----\n")# 循环读取路径下的文件并操作
for video_name in dirs:outputPath = "F:\\testkk\\images\\"+video_name[:-4] + "\\"# os.mkdir(outputPath)print("start\n")print(videos_src_path + video_name)vc = cv2.VideoCapture(videos_src_path + video_name)# 初始化,并读取第一帧# rval表示是否成功获取帧# frame是捕获到的图像rval, frame = vc.read()# 获取视频fpsfps = vc.get(cv2.CAP_PROP_FPS)# 获取每个视频帧数frame_all = vc.get(cv2.CAP_PROP_FRAME_COUNT)print("[INFO] 视频FPS: {}".format(fps))print("[INFO] 视频总帧数: {}".format(frame_all))# 每隔n帧保存一张图片frame_interval = 30# 统计当前帧frame_count = 1# count=0while rval:rval, frame = vc.read()# 隔n帧保存一张图片if frame_count % frame_interval == 0:# 当前帧不为None，能读取到图片时if frame is not None:filename = outputPath + "Kidney_tumors_{}.jpg".format(count)# 水平、垂直翻转frame = cv2.flip(frame, 0)frame = cv2.flip(frame, 1)# 旋转180°frame = np.rot90(frame)frame = np.rot90(frame)cv2.imwrite(filename, frame)count += 1print("保存图片:{}".format(filename))frame_count += 1# 将成功抽帧的视频名称写入txt文件，方便检查file = open(f, "a")file.write(video_name + "\n")# 关闭视频文件vc.release()print("[INFO] 总共保存：{}张图片\n".format(count))

遇到问题：
global cap_ffmpeg_impl.hpp:1541 grabFrame packet read max attempts exceeded, if your video have multiple streams (video, audio) try to increase attempt limit by setting environment variable OPENCV_FFMPEG_READ_ATTEMPTS (current value is 4096)
解决方法：
Windows：给OPENCV_FFMPEG_READ_ATTEMPTS 设置一个和视频大小一样的value值，然后重启电脑，再来执行，就解决了。
Linux：终端执行
$ export OPENCV_FFMPEG_READ_ATTEMPTS=4001989068

抽出来的图片，接下来用YOLO的目标检测模型预测，得到坐标文件

from ultralytics import YOLO# 读取模型，这里传入训练好的模型
model = YOLO('yolov8m.pt')# 模型预测，save=True 的时候表示直接保存yolov8的预测结果
metrics = model.predict(source='F:\\testkk\\images\\GH040001',imgsz=640,project='runs/detect',save=True)

在这里插入图片描述
【预测结果已经还OK了，几乎不用咋修改标注文件了】

然后：【txt转json】

import os
import json
import base64
import cv2def read_txt_file(txt_file):with open(txt_file, 'r') as f:lines = f.readlines()data = []for line in lines:line = line.strip().split()class_name = line[0]bbox = [coord for coord in line[1:]]data.append({'class_name': class_name, 'bbox': bbox})return datadef convert_to_labelme(data, image_path, image_size):labelme_data = {'version': '4.5.6','flags': {},'shapes': [],'imagePath': json_image_path,'imageData': None,'imageHeight': image_size[0],'imageWidth': image_size[1]}for obj in data:dx = obj['bbox'][0]dy = obj['bbox'][1]dw = obj['bbox'][2]dh = obj['bbox'][3]w = eval(dw) * image_size[1]h = eval(dh) * image_size[0]center_x = eval(dx) * image_size[1]center_y = eval(dy) * image_size[0]x1 = center_x - w/2y1 = center_y - h/2x2 = center_x + w/2y2 = center_y + h/2# x1 = eval(obj['bbox'][0]) * image_size[1]# y1 = eval(obj['bbox'][1]) * image_size[0]# x2 = eval(obj['bbox'][2]) * image_size[1]# y2 = eval(obj['bbox'][3]) * image_size[0]if obj['class_name'] == '0': #判断对应的标签名称，写入json文件中label = str('person')elif obj['class_name'] == '2':label = str('car')else:continueshape_data = {'label': label,'points': [[x1, y1], [x2, y2]],'group_id': None,'shape_type': 'rectangle','flags': {}}labelme_data['shapes'].append(shape_data)return labelme_datadef save_labelme_json(labelme_data, image_path, output_file):with open(image_path, 'rb') as f:image_data = f.read()labelme_data['imageData'] = base64.b64encode(image_data).decode('utf-8')with open(output_file, 'w') as f:json.dump(labelme_data, f, indent=4)# 设置文件夹路径和输出文件夹路径
txt_folder = "D:\\yoloProject\\ultralytics-registry\\runs\\detect\\predict5\\labels"  # 存放LabelImg标注的txt文件的文件夹路径
output_folder = "F:\\testkk\\images\\GH040001_json"  # 输出LabelMe标注的json文件的文件夹路径
img_folder = "F:\\testkk\\images\\GH040001" #存放对应标签的图片文件夹路径# 创建输出文件夹
if not os.path.exists(output_folder):os.makedirs(output_folder)# 遍历txt文件夹中的所有文件
for filename in os.listdir(txt_folder):if filename.endswith('.txt'):# 生成对应的输出文件名output_filename = os.path.splitext(filename)[0] + '.json'# 读取txt文件txt_file = os.path.join(txt_folder, filename)data = read_txt_file(txt_file)# 设置图片路径和尺寸image_filename = os.path.splitext(filename)[0] + '.jpg'  # 图片文件名与txt文件名相同，后缀为.jpgimage_path = os.path.join(img_folder, image_filename)# image_size = (1280, 720)  # 根据实际情况修改json_image_path = image_path.split('\\')[-1]image_size = cv2.imread(image_path).shape# 转化为LabelMe格式labelme_data = convert_to_labelme(data, image_path, image_size)# 保存为LabelMe JSON文件output_file = os.path.join(output_folder, output_filename)save_labelme_json(labelme_data, image_path, output_file)