小工具之视频抽帧

'''
视频抽帧工具，所有视频所在目录以及抽帧图片保存路径
单个视频抽帧操作步骤：
选择文件路径->选择保存路径->拖动跳帧间隔->点击抽取帧
批量视频抽帧操作步骤：
选择文件夹路径->选择保存路径->拖动跳帧间隔->点击抽取帧
'''


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_())