python
复制代码
'''
视频抽帧工具，所有视频所在目录以及抽帧图片保存路径
单个视频抽帧操作步骤：
选择文件路径->选择保存路径->拖动跳帧间隔->点击抽取帧
批量视频抽帧操作步骤：
选择文件夹路径->选择保存路径->拖动跳帧间隔->点击抽取帧
'''
import sys
import cv2
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QPushButton, QLabel, QFileDialog, QSlider, QHBoxLayout, QProgressBar
from PyQt5.QtCore import Qt
import os
from shortuuid import uuid
# include image suffixes
IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp', 'pfm'
# include video suffixes
VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv'
def is_img(file_path: str):
return file_path.split('.')[-1].lower() in IMG_FORMATS
def is_video(file_path: str):
return file_path.split('.')[-1].lower() in VID_FORMATS
class VideoFrameExtractor(QMainWindow):
def __init__(self):
super().__init__()
self.videoCapture = None
self.frame_counter = 0
self.save_path = ""
self.central_widget = QWidget(self)
self.setCentralWidget(self.central_widget)
self.layout = QVBoxLayout()
self.central_widget.setLayout(self.layout)
self.progress_bar = QProgressBar()
self.layout.addWidget(self.progress_bar)
self.video_label = QLabel(self)
self.layout.addWidget(self.video_label)
file_select_layout = QHBoxLayout()
self.file_button = QPushButton("选择文件路径", self)
self.file_button.clicked.connect(self.openFile)
file_select_layout.addWidget(self.file_button)
self.dir_button = QPushButton("选择文件夹路径", self)
self.dir_button.clicked.connect(self.openFileDir)
file_select_layout.addWidget(self.dir_button)
self.layout.addLayout(file_select_layout)
save_start_layout = QHBoxLayout()
self.save_button = QPushButton("选择保存路径", self)
self.save_button.clicked.connect(self.openSavePath)
save_start_layout.addWidget(self.save_button)
self.extract_button = QPushButton("抽取帧", self)
self.extract_button.clicked.connect(self.extractFrames)
self.extract_button.setEnabled(False)
save_start_layout.addWidget(self.extract_button)
self.layout.addLayout(save_start_layout)
param_start_layout = QHBoxLayout()
self.frame_label = QLabel("跳帧间隔: 1", self)
param_start_layout.addWidget(self.frame_label)
self.frame_slider = QSlider(Qt.Horizontal, self)
self.frame_slider.setRange(1, 100)
self.frame_slider.setValue(1)
self.frame_slider.valueChanged.connect(self.update_skip_param)
param_start_layout.addWidget(self.frame_slider)
self.layout.addLayout(param_start_layout)
self.videos = []
self.setGeometry(100, 100, 800, 100)
self.setWindowTitle("视频抽帧应用")
self.show()
def update_skip_param(self):
self.skip_param = self.frame_slider.value()
self.frame_label.setText(f"跳帧间隔: {self.skip_param}")
def openFileDir(self):
options = QFileDialog.Options()
# options |= QFileDialog.DontUseNativeDialog
options |= QFileDialog.ShowDirsOnly
file_dir = QFileDialog.getExistingDirectory(
self, "选择视频文件夹", "", options=options)
if file_dir:
self.extract_button.setEnabled(True)
all_files = os.listdir(file_dir)
self.videos = [os.path.join(file_dir, file)
for file in all_files if is_video(file)]
def openFile(self):
options = QFileDialog.Options()
options |= QFileDialog.ReadOnly
file_path, _ = QFileDialog.getOpenFileName(
self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mkv)", options=options)
if file_path:
self.extract_button.setEnabled(True)
self.videos = [file_path]
def openSavePath(self):
options = QFileDialog.Options()
options |= QFileDialog.ReadOnly
directory = QFileDialog.getExistingDirectory(
self, "选择保存路径", options=options)
if directory:
self.save_path = directory
def disable_elems(self):
self.extract_button.setEnabled(False)
self.file_button.setEnabled(False)
self.dir_button.setEnabled(False)
self.save_button.setEnabled(False)
self.frame_slider.setEnabled(False)
def enable_elems(self):
self.extract_button.setEnabled(True)
self.file_button.setEnabled(True)
self.dir_button.setEnabled(True)
self.save_button.setEnabled(True)
self.frame_slider.setEnabled(True)
def handle(self):
pass
def extractFrames(self):
self.disable_elems()
frame_skip = self.frame_slider.value()
max_value = len(self.videos)
self.progress_bar.setValue(0)
for idx, video_path in enumerate(self.videos):
cap = cv2.VideoCapture(video_path)
prefix = str(uuid())
# 以视频名称命名文件夹
save_name = os.path.join(os.path.splitext(
os.path.basename(video_path))[0])
save_dir = os.path.join(self.save_path, save_name)
os.makedirs(save_dir, exist_ok=True)
frame_idx = -1
while True:
frame_idx += 1
cap.grab()
if frame_idx % frame_skip != 0:
continue
ret, frame = cap.retrieve()
if not ret:
break
cv2.imencode('.jpg', frame, [int(cv2.IMWRITE_JPEG_QUALITY), 100])[
1].tofile(f"{save_dir}/{prefix}_{str(frame_idx).zfill(6)}.jpg")
cap.release()
self.progress_bar.setValue(int((idx+1) / max_value * 100))
QApplication.processEvents() # 实时刷新界面，防止界面卡住
self.progress_bar.repaint()
# self.progress_bar.setFormat(f"{idx}% 完成")
self.enable_elems()
if __name__ == "__main__":
app = QApplication(sys.argv)
window = VideoFrameExtractor()
sys.exit(app.exec_())
