import os, shutil
from docarray import Document, DocumentArray

# https://blog.csdn.net/Jina_AI/article/details/128475707

class Screenshot:
    # 视频名称
    videoName = ''
    # 视频所在的 父级路径
    videoPath = '../../source/'
    # 视频帧截图后输出的 父级路径
    outputKeyframes = '../../dist/'
    d = ''
    keyframes = DocumentArray()

    def __init__(self, video_name):
        self.videoName = video_name
        self.d = Document(uri=self.videoPath + self.videoName).load_uri_to_video_tensor(only_keyframes=False)
        self.createdDir()
        self.start()

    def start(self):
        for i in range(len(self.d.tensor)):
            if i in self.d.tags['keyframe_indices']:
                keyframe = Document(
                    tensor=self.d.tensor[i], tags={'index': len(self.keyframes)}
                )
                keyframe.save_image_tensor_to_file(file=f'{self.outputKeyframes}/{self.videoName}/{len(self.keyframes)}.png')
                self.keyframes.append(keyframe)


    def createdDir(self):
        if not os.path.isdir(self.outputKeyframes + self.videoName):
            os.makedirs(self.outputKeyframes + self.videoName, exist_ok=True)


Screenshot('yjkon.mp4')