Implement ECC, webcam capture

1 files changed, 67 insertions, 17 deletions

diff --git a/encoder.py b/encoder.py
index c5d76e1..6fe05c4 100644
--- a/encoder.py
+++ b/encoder.py
@@ -1,14 +1,38 @@
+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
 
-fps = 30
-h = 200
-w = 200
-c = 0
+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()
+
+
+cheight = args.height // 10
+cwidth = args.width // 10
+midwidth = args.width - 2 * cwidth
+frame_size = args.height * args.width - 4 * cheight * cwidth
+# 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):
@@ -16,7 +40,8 @@ class EncoderWidget(QWidget):
         super().__init__()
         self.timer = QTimer(self)
         self.timer.timeout.connect(self.update)
-        self.timer.start(1000 // fps)
+        self.timer.start(1000 // args.fps)
+        self.idx = 0
         self.label = QLabel(self)
         layout = QVBoxLayout(self)
         layout.addWidget(self.label)
@@ -25,20 +50,45 @@ class EncoderWidget(QWidget):
         self.showFullScreen()
 
     def update(self):
-        global c
-        if c == 0:
-            array = np.random.randint(0, 256, (h, w, 3))
-            c = 1
-        else:
-            array = np.zeros((h, w, 3))
-            c = 0
-        img = Image.fromarray(array.astype(np.uint8), mode="RGB")
+        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)))
+        self.idx = (self.idx + 1) % len(packets)
+        frame = np.concatenate(
+            (
+                np.concatenate(
+                    (
+                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                        frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
+                        np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
+                    ),
+                    axis=1,
+                ),
+                frame_data[
+                    cheight * midwidth : frame_size - cheight * midwidth
+                ].reshape((args.height - 2 * cheight, args.width)),
+                np.concatenate(
+                    (
+                        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
+                    ),
+                    axis=1,
+                ),
+            )
+        )
+        color_frame = np.stack(
+            ((frame & 0x03) << 6, (frame >> 2 & 0x07) << 5, (frame >> 5 & 0x7) << 5),
+            axis=-1,
+        )
+        img = Image.fromarray(color_frame)
         qt_img = ImageQt.ImageQt(img)
         pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
         self.label.setPixmap(pixmap)
 
 
-if __name__ == "__main__":
-    app = QApplication(sys.argv)
-    widget = EncoderWidget()
-    sys.exit(app.exec())
+app = QApplication([])
+widget = EncoderWidget()
+sys.exit(app.exec())