File size: 1,181 Bytes
ce0d4fb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import cv2


def video_reader(file_path, targetFPS=9, targetWidth=None, to_rgb=True, prompt=None):
    cap = cv2.VideoCapture(file_path)
    sourceFPS = int(cap.get(cv2.CAP_PROP_FPS))

    if sourceFPS < targetFPS:
        raise ValueError("sourceFPS < targetFPS: {} < {}".format(sourceFPS, targetFPS))

    fpsDiv = 3  # sourceFPS // targetFPS
    print("sourceFPS: {}, targetFPS: {}, fpsDiv: {}".format(sourceFPS, targetFPS, fpsDiv))

    frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frameWidth = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frameHeight = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    print("frameCount: {}, frameWidth: {}, fpsDiv: {}".format(frameCount, frameWidth, frameHeight))

    if targetWidth:
        targetHeight = int(targetWidth * frameHeight / frameWidth)

    fc = 0
    ret = True

    while fc < frameCount and ret:
        ret, img = cap.read()
        if fc % fpsDiv == 0:
            if targetWidth:
                img = cv2.resize(img, (targetWidth, targetHeight), interpolation=cv2.INTER_AREA)
            if to_rgb:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            yield img
        fc += 1

    cap.release()