aboutsummaryrefslogtreecommitdiff
path: root/encoder.py
blob: b20cd5696085a88885b538762eb6719086b20fa3 (plain)
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import argparse
import sys
import cv2
import numpy as np
from creedsolo import RSCodec
from PyQt6.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout
from PyQt6.QtGui import QPixmap
from PyQt6.QtCore import QTimer
from PIL import Image, ImageQt
from raptorq import Encoder

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("file", help="input file")
parser.add_argument("--height", help="grid height", default=100, type=int)
parser.add_argument("--width", help="grid width", default=100, type=int)
parser.add_argument("--fps", help="framerate", default=30, type=int)
parser.add_argument("--level", help="error correction level", default=0.1, type=float)
parser.add_argument("--video", help="output file for encoded video")
parser.add_argument("--scale", help="scale of new frames", default=2, type=int)
args = parser.parse_args()


if args.video:
    cap = cv2.VideoCapture(args.file)
    args.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) / args.scale)
    args.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) / args.scale)

# Make corners
cheight = cwidth = max(args.height // 10, args.width // 10)
wcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b11111111), ((0, 1), (0, 1)))
rcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b00000111), ((0, 1), (1, 0)))
gcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b00111000), ((1, 0), (0, 1)))
bcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b11000000), ((1, 0), (1, 0)))
midwidth = args.width - 2 * cwidth
frame_size = args.height * args.width - 4 * cheight * cwidth
frame_xor = np.arange(frame_size, dtype=np.uint8)
# reedsolo breaks message into 255-byte chunks
# raptorq can add up to 4 extra bytes
rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4

with open(args.file, "rb") as f:
    data = f.read()
rsc = RSCodec(int(args.level * 255))
encoder = Encoder.with_defaults(data, rs_size)
packets = encoder.get_encoded_packets(int(len(data) / rs_size * args.level))

print("Data length:", len(data))
print("Packets:", len(packets))

idx = 0


def get_frame():
    global idx
    frame_data = np.array(rsc.encode(packets[idx]))
    # Pad frame to fit frame_size since raptorq might not add 4 bytes
    frame_data = np.pad(frame_data, (0, frame_size - len(frame_data))) ^ frame_xor
    idx = (idx + 1) % len(packets)
    frame = np.concatenate(
        (
            np.concatenate(
                (
                    wcorner,
                    frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
                    rcorner,
                ),
                axis=1,
            ),
            frame_data[cheight * midwidth : frame_size - cheight * midwidth].reshape(
                (args.height - 2 * cheight, args.width)
            ),
            np.concatenate(
                (
                    gcorner,
                    frame_data[frame_size - cheight * midwidth :].reshape(
                        (cheight, midwidth)
                    ),
                    bcorner,
                ),
                axis=1,
            ),
        )
    )
    return np.stack(
        (
            (frame & 0b00000111) * 255 // 7,
            (frame >> 3 & 0b00000111) * 255 // 7,
            (frame >> 6 & 0b00000011) * 255 // 3,
        ),
        axis=-1,
    ).astype(np.uint8)


class EncoderWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update)
        self.timer.start(1000 // args.fps)
        self.label = QLabel(self)
        layout = QVBoxLayout(self)
        layout.addWidget(self.label)
        layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(layout)
        self.showFullScreen()

    def update(self):
        img = Image.fromarray(get_frame())
        qt_img = ImageQt.ImageQt(img)
        pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
        self.label.setPixmap(pixmap)


if args.video:
    out = cv2.VideoWriter(
        args.video,
        cv2.VideoWriter_fourcc(*"mp4v"),
        cap.get(cv2.CAP_PROP_FPS),
        (args.scale * args.width, args.scale * args.height),
    )
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame[: args.scale * cheight, : args.scale * cwidth] = 255
        frame[: args.scale * cheight, args.scale * (args.width - cwidth) :] = 255
        frame[args.scale * (args.height - cheight) :, : args.scale * cwidth] = 255
        frame[
            args.scale * (args.height - cheight) :, args.scale * (args.width - cwidth) :
        ] = 255
        out.write(
            (
                frame.astype(np.int64)
                * np.repeat(np.repeat(get_frame(), args.scale, 0), args.scale, 1)
                / 255
            ).astype(np.uint8)
        )
else:
    input("Seizure warning!")
    app = QApplication([])
    widget = EncoderWidget()
    sys.exit(app.exec())