aboutsummaryrefslogtreecommitdiff
path: root/encoder.py
blob: a4ef7c19dfa5875d262a3948241566aec9e18300 (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
import argparse
import sys
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="output file for decoded data")
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)
args = parser.parse_args()


# 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))
input("Seizure warning!")


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.idx = 0
        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):
        frame_data = np.array(rsc.encode(packets[self.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
        self.idx = (self.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,
                ),
            )
        )
        color_frame = np.stack(
            (
                (frame & 0b00000111) * 255 // 7,
                (frame >> 3 & 0b00000111) * 255 // 7,
                (frame >> 6 & 0b00000011) * 255 // 3,
            ),
            axis=-1,
        )
        img = Image.fromarray(color_frame.astype(np.uint8))
        qt_img = ImageQt.ImageQt(img)
        pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
        self.label.setPixmap(pixmap)


app = QApplication([])
widget = EncoderWidget()
sys.exit(app.exec())