aboutsummaryrefslogtreecommitdiff
path: root/decoder.py
blob: e58e246ba9df5c9e69ea2b997517a4f788a4cc11 (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
import argparse
import cv2
import numpy as np
from creedsolo import RSCodec
from raptorq import Decoder

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("file", help="output file for decoded data")
parser.add_argument("--len", help="number of bytes to decode", default=2**16, type=int)
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("--device", help="camera device index", default=1, type=int)
args = parser.parse_args()

cheight = args.height // 10
cwidth = args.width // 10
frame_size = args.height * args.width - 4 * cheight * cwidth
frame_xor = np.arange(frame_size, dtype=np.uint8)
rs_size = int(frame_size * (1 - args.level))

rsc = RSCodec(frame_size - rs_size)
raptor_decoder = Decoder.with_defaults(args.len, rs_size)
data = None
cap = cv2.VideoCapture(args.device)
while data is None:
    ret, raw_frame = cap.read()
    if not ret:
        continue
    color_frame = decode(raw_frame)  # TODO
    frame = (
        (color_frame[:, :, 0] >> 5 & 0b00000111)
        + (color_frame[:, :, 1] >> 2 & 0b00111000)
        + (color_frame[:, :, 2] & 0b11000000)
    )
    frame_data = np.concatenate(
        (
            frame[:cheight, cwidth : args.width - cwidth].flatten(),
            frame[cheight : args.height - cheight].flatten(),
            frame[args.heigth - cheight, cwidth : args.width - cwidth].flatten(),
        )
    )
    data = raptor_decoder.decode(rsc.decode(frame_data ^ frame_xor))
with open(args.file, "wb") as f:
    f.write(data)
cap.release()