Move find_corner to global namespace, use erase_pos for better ECC

2 files changed, 20 insertions, 20 deletions

diff --git a/decoder.py b/decoder.py
index dccb24e..2467328 100644
--- a/decoder.py
+++ b/decoder.py
@@ -1,5 +1,4 @@
 import argparse
-import collections
 import sys
 import traceback
 import cv2
@@ -26,6 +25,19 @@ rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 25
 rsc = RSCodec(int(args.level * 255))
 decoder = Decoder.with_defaults(args.size, rs_size)
 
+
+def find_corner(A, f):
+    cx, cy = A.shape[:2]
+    # Resize so smaller dim is 5
+    scale = min(cx // 5, cy // 5)
+    B = cv2.resize(A, (cy // scale, cx // scale), interpolation=cv2.INTER_AREA)
+    guess = np.array(np.unravel_index(np.argmax(f(B)), B.shape[:2])) * scale + scale // 2
+    mask = cv2.floodFill(A, np.empty(0), tuple(reversed(guess)), 1, 10, 10, cv2.FLOODFILL_MASK_ONLY)[2][
+        1:-1, 1:-1
+    ].astype(bool)
+    return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64)
+
+
 if args.input.isdecimal():
     args.input = int(args.input)
 cap = cv2.VideoCapture(args.input)
@@ -36,24 +48,13 @@ while data is None:
         if not ret:
             print("End of stream")
             sys.exit(1)
-        cv2.imshow("", raw_frame)
-        cv2.waitKey(33)
+        # cv2.imshow("", raw_frame)
+        # cv2.waitKey(33)
         # raw_frame is a uint8 BE CAREFUL
         raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)
 
         X, Y = raw_frame.shape[:2]
         cx, cy = X // 4, Y // 4
-        scale = min(cx // 5, cy // 5)
-        # Resize so smaller dim is 5
-
-        def find_corner(A, f):
-            B = cv2.resize(A, (cy // scale, cx // scale), interpolation=cv2.INTER_AREA)
-            guess = np.array(np.unravel_index(np.argmax(f(B)), B.shape[:2])) * scale + scale // 2
-            mask = cv2.floodFill(A, np.empty(0), tuple(reversed(guess)), 1, 10, 10, cv2.FLOODFILL_MASK_ONLY)[2][
-                1:-1, 1:-1
-            ].astype(bool)
-            return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64)
-
         widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: (np.std(B, axis=2) < 35) * np.sum(B, axis=2))
         ridx, rcol = find_corner(raw_frame[:cx, Y - cy :], lambda B: B[:, :, 0] - B[:, :, 1] - B[:, :, 2])
         ridx[1] += Y - cy
@@ -114,11 +115,10 @@ while data is None:
                 frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
             )
         )
-        frame = ((frame[::2] << 4) + frame[1::2]) ^ frame_xor
+        frame = (frame[::2] << 4) + frame[1::2]
         frame = np.pad(frame, (0, (len(frame) + 254) // 255 * 255 - len(frame)))
         frame = np.ravel(frame.reshape(255, len(frame) // 255), "F")[: frame_size // 2]
-
-        data = decoder.decode(bytes(rsc.decode(frame)[0]))
+        data = decoder.decode(bytes(rsc.decode(frame ^ frame_xor, erase_pos=list(np.where(frame == 0)[0]))[0]))
         print("Decoded frame")
     except KeyboardInterrupt:
         sys.exit()
diff --git a/encoder.py b/encoder.py
index 0c68866..3b67cf0 100644
--- a/encoder.py
+++ b/encoder.py
@@ -33,8 +33,7 @@ with open(args.input, "rb") as f:
     data = f.read()
 rsc = RSCodec(int(args.level * 255))
 encoder = Encoder.with_defaults(data, rs_size)
-# This formula is wrong 🤷
-packets = encoder.get_encoded_packets(int(len(data) / rs_size * args.level))
+packets = encoder.get_encoded_packets(int(len(data) / rs_size * (1 / (1 - args.level) - 1)))
 
 # Make corners
 wcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b1111), ((0, 1), (0, 1)))
@@ -52,10 +51,11 @@ def get_frame():
     global idx
     frame = np.array(rsc.encode(packets[idx]))
     idx = (idx + 1) % len(packets)
+    frame = np.pad(frame, (0, frame_size // 2 - len(frame))) ^ frame_xor
     # Add 4 bytes, pad frame to be multiple of 255
     frame = np.pad(frame, (0, (len(frame) + 258) // 255 * 255 - len(frame)))
     # Space out elements in each size 255 chunk
-    frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2] ^ frame_xor
+    frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2]
     frame = np.ravel(np.column_stack((frame >> 4, frame & 0b1111)))
     frame = np.concatenate(
         (