XOR frames with arange, implement color to 8-bit in decoder, add camera device index flag

2 files changed, 22 insertions, 14 deletions

diff --git a/decoder.py b/decoder.py
index 357db39..e58e246 100644
--- a/decoder.py
+++ b/decoder.py
@@ -11,30 +11,37 @@ 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.size, rs_size)
+raptor_decoder = Decoder.with_defaults(args.len, rs_size)
 data = None
-cap = cv2.VideoCapture(0)
+cap = cv2.VideoCapture(args.device)
 while data is None:
-    ret, frame = cap.read()
+    ret, raw_frame = cap.read()
     if not ret:
         continue
-    frame_full = decode(frame)  # TODO
+    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_full[:cheight, cwidth : args.width - cwidth].flatten(),
-            frame_full[cheight : args.height - cheight].flatten(),
-            frame_full[args.heigth - cheight, cwidth : args.width - cwidth].flatten(),
+            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))
+    data = raptor_decoder.decode(rsc.decode(frame_data ^ frame_xor))
 with open(args.file, "wb") as f:
     f.write(data)
 cap.release()
diff --git a/encoder.py b/encoder.py
index 6fe05c4..cc4d7cb 100644
--- a/encoder.py
+++ b/encoder.py
@@ -21,6 +21,7 @@ cheight = args.height // 10
 cwidth = args.width // 10
 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
@@ -52,15 +53,15 @@ class EncoderWidget(QWidget):
     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_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(
                     (
-                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                         frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
-                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                     ),
                     axis=1,
                 ),
@@ -69,18 +70,18 @@ class EncoderWidget(QWidget):
                 ].reshape((args.height - 2 * cheight, args.width)),
                 np.concatenate(
                     (
-                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                         frame_data[frame_size - cheight * midwidth :].reshape(
                             (cheight, midwidth)
                         ),
-                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                     ),
                     axis=1,
                 ),
             )
         )
         color_frame = np.stack(
-            ((frame & 0x03) << 6, (frame >> 2 & 0x07) << 5, (frame >> 5 & 0x7) << 5),
+            (frame << 5 & 0b11100000, frame << 2 & 0b11100000, frame & 0b11000000),
             axis=-1,
         )
         img = Image.fromarray(color_frame)