104AFinal/process.py at main · nhamovitz/104AFinal · GitHub

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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97


import cv2
import numpy as np

def sparse(path: str, cut_proportion = 2):
    cap = cv2.VideoCapture(path)

    frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    frame_rate = cap.get(cv2.CAP_PROP_FPS)

    print(f"Making sparse vid; original shape {(frame_count, height, width, 3)}.")

    # these come as floats for some reason
    # note: we want to round up, not round or truncate
    # that's weird to me, but you get an off-by-one error with n = 3 if this is `int(frame_count // cut_proportion), and same with n = 10 and `round(frame_count / cut_proportion)`
    keep_count = None
    if frame_count % cut_proportion == 0:
        keep_count = int(frame_count // cut_proportion)
    else:
        keep_count = int(frame_count // cut_proportion + 1)
    # print(frame_count, int(frame_count // cut_proportion), frame_count / cut_proportion, keep_count)
    width = int(width)
    height = int(height)

    sparse_vid = np.empty((keep_count, height, width, 3), dtype=np.uint8)

    kept_frames = []
    frame_number = 0

    while (cap.isOpened()):
        _, frame = cap.read()

        if frame_number % cut_proportion == 0:
            index = round(frame_number / cut_proportion)
            sparse_vid[index, :, :, :] = frame
            kept_frames.append(frame_number)

        if frame_number == frame_count - 1:
            _, frame = cap.read()
            frame_number += 1
            assert frame is None
            cap.release()
            break

        frame_number += 1

    print(f"Kept every {cut_proportion}; {len(kept_frames)} frames total; indices {kept_frames}.")
    return sparse_vid, kept_frames, frame_rate


def write_black_with_codec(codec, extension):
    height, width = 720, 1280
    fourcc = cv2.VideoWriter_fourcc(*codec)
    writer = cv2.VideoWriter(f".\\media\\test\\{codec}_{extension}.{extension}", fourcc, 30, (height, width))
    for _ in range(150):
        blank_frame = np.zeros((height, width))
        writer.write(blank_frame)
    writer.release()

def test_writing_video():
    for codec in ('DIVX', 'XVID', 'MJPG', 'X264', 'WMV1', 'WMV2', 'MP4V', 'MPEG', 'H264', 'mp4v', 'mpv4'):
        for ext in ('mp4', 'avi', 'mpg'):
            try:
                write_black_with_codec(codec, ext)
            except Exception as e:
                print(f"{codec=}, {ext=}, {e=}")


def run_demo():
    from pathlib import Path
    demo = str(Path('.') / 'media' / 'keys.mp4')

    # sparse_vid, kept, _ = sparse(demo)
    # print("every 2", sparse_vid.shape, kept, len(kept))

    # sparse_vid, kept, _ = sparse(demo, 3)
    # print("every 3", sparse_vid.shape, kept, len(kept))

    # sparse_vid, kept, _ = sparse(demo, 10)
    # print("every 10", sparse_vid.shape, kept, len(kept))

    # redone_vid, _, _ = sparse(demo, 1)

    # print(redone_vid.shape)

    interval = 5
    sparse_vid, kept, frame_rate = sparse(demo, interval)
    print(f"every {interval}", sparse_vid.shape, kept, len(kept))

    return sparse_vid, kept

if __name__ == '__main__':
    run_demo()