在做室外语义分割、视觉导航与定位的时候,通常会用对一个连续的视频帧进行测试,除去常用数据集外,也经常会自己制作一些数据集,这个工具类可以按需求对视频进行分帧提取,封装好了直接可以使用,根据需求更改视频路径,保存路径和帧率即可。
import os
import cv2
import shutil
VIDEO_PATH = 'test.mp4' # 视频地址
EXTRACT_FOLDER = 'video' # 存放帧图片的位置
EXTRACT_FREQUENCY = 1 # 帧提取频率
def extract_frames(video_path, dst_folder, index):
# 主操作
video = cv2.VideoCapture()
if not video.open(video_path):
print("can not open the video")
exit(1)
count = 1
while True:
_, frame = video.read()
if frame is None:
break
if count % EXTRACT_FREQUENCY == 0:
save_path = "{}/{:>03d}.jpg".format(dst_folder, index)
cv2.imwrite(save_path, frame)
index += 1
count += 1
video.release()
# 打印出所提取帧的总数
print("Totally save {:d} pics".format(index - 1))
def main():
# 递归删除之前存放帧图片的文件夹,并新建一个
try:
shutil.rmtree(EXTRACT_FOLDER)
except OSError:
pass
if not os.path.exists(EXTRACT_FOLDER):
os.mkdir(EXTRACT_FOLDER)
# 抽取帧图片,并保存到指定路径
extract_frames(VIDEO_PATH, EXTRACT_FOLDER, 1)
if __name__ == '__main__':
main()
另外再po一个工具类,用来获取指定路径下的文件夹或者文件名,因为我在使用网络进行验证的时候,要对多个数据集进行验证,一个一个图片去预测验证的时候太慢了,所以就获取整个文件夹的名称进行预测。
def getFiles(rootDir):
files = []
for lists in os.listdir(rootDir):
path = os.path.join(rootDir, lists)
files.append(path)
# print(path)
if os.path.isdir(path):
getFiles(path)
return files