高效率制作数据集【按这个流程走,速度很顶】
本次制作,1059张图片【马路上流动车辆】
几乎就是全自动了,只要视频拍得好,YOLO辅助制作数据集就效率极高
视频中的图片抽取:
【由于视频内存过大,遇到报错执行失败,解决方法已附加在代码下面】
import cv2
import os
import pdb
import numpy as np
#from glob2 import glob
videos_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()
# 获取视频fps
fps = 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=0
while 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 += 1
print("保存图片:{}".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)
解决方法:
给OPENCV_FFMPEG_READ_ATTEMPTS 设置一个和视频大小一样的value值,然后重启电脑,再来执行,就解决了。
抽出来的图片,接下来用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 cv2
def 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 data
def 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/2
y1 = center_y - h/2
x2 = center_x + w/2
y2 = 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:
continue
shape_data = {
'label': label,
'points': [[x1, y1], [x2, y2]],
'group_id': None,
'shape_type': 'rectangle',
'flags': {}
}
labelme_data['shapes'].append(shape_data)
return labelme_data
def 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文件名相同,后缀为.jpg
image_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)
最后:修改OK后,再把JSON转TXT,作为样本数据集:
https://blog.csdn.net/weixin_43624549/article/details/139532142