Downsample+flood fill algorithm (it almost works!)

1 files changed, 17 insertions, 8 deletions

diff --git a/encoder.py b/encoder.py
index cc4d7cb..a4ef7c1 100644
--- a/encoder.py
+++ b/encoder.py
@@ -17,8 +17,12 @@ parser.add_argument("--level", help="error correction level", default=0.1, type=
 args = parser.parse_args()
 
 
-cheight = args.height // 10
-cwidth = args.width // 10
+# 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)
@@ -31,6 +35,7 @@ with open(args.file, "rb") as f:
 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!")
@@ -59,9 +64,9 @@ class EncoderWidget(QWidget):
             (
                 np.concatenate(
                     (
-                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
+                        wcorner,
                         frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
-                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
+                        rcorner,
                     ),
                     axis=1,
                 ),
@@ -70,21 +75,25 @@ class EncoderWidget(QWidget):
                 ].reshape((args.height - 2 * cheight, args.width)),
                 np.concatenate(
                     (
-                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
+                        gcorner,
                         frame_data[frame_size - cheight * midwidth :].reshape(
                             (cheight, midwidth)
                         ),
-                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
+                        bcorner,
                     ),
                     axis=1,
                 ),
             )
         )
         color_frame = np.stack(
-            (frame << 5 & 0b11100000, frame << 2 & 0b11100000, frame & 0b11000000),
+            (
+                (frame & 0b00000111) * 255 // 7,
+                (frame >> 3 & 0b00000111) * 255 // 7,
+                (frame >> 6 & 0b00000011) * 255 // 3,
+            ),
             axis=-1,
         )
-        img = Image.fromarray(color_frame)
+        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)