4-bit color for encoder, fix a ton of bugs

3 files changed, 44 insertions, 74 deletions

diff --git a/README.md b/README.md
index 42fac66..c2ac901 100644
--- a/README.md
+++ b/README.md
@@ -3,8 +3,5 @@
 Formatting: `black -l 120 *.py`
 
 TODO:
-- Only search corner for corner
-- Fast BFS using numpy (OpenCV floodfill?)
 - Limit BFS time
-- Spread out RS bytes (still need to impl in decoder)
-- 4bit color (still need to impl in decoder)
+- Write better ECC
diff --git a/decoder.py b/decoder.py
index 860d585..dccb24e 100644
--- a/decoder.py
+++ b/decoder.py
@@ -1,6 +1,7 @@
 import argparse
 import collections
 import sys
+import traceback
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
@@ -19,8 +20,8 @@ args = parser.parse_args()
 
 cheight = cwidth = max(args.height // 10, args.width // 10)
 frame_size = args.height * args.width - 4 * cheight * cwidth
-frame_xor = np.arange(frame_size, dtype=np.uint8)
-rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4
+frame_xor = np.arange(frame_size // 2, dtype=np.uint8)
+rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
 
 rsc = RSCodec(int(args.level * 255))
 decoder = Decoder.with_defaults(args.size, rs_size)
@@ -35,65 +36,32 @@ while data is None:
         if not ret:
             print("End of stream")
             sys.exit(1)
-        raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB).astype(np.float64)
+        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]
-        scale = min(X // 20, Y // 20)
-        # Resize so smaller dim is 20
-        # Use fast default interpolation for factor of 4
-        # Then switch to good slow interpolation
-        dframe = cv2.resize(
-            cv2.resize(raw_frame, (Y // 4, X // 4)),
-            (Y // scale, X // scale),  # OpenCV swaps them
-            interpolation=cv2.INTER_AREA,
-        )
-        # plt.imshow(dframe.astype(np.uint8))
-        # plt.show()
-
-        def max_in_orig(x):
-            return tuple(np.array(np.unravel_index(np.argmax(x), x.shape)) * scale + scale // 2)
+        cx, cy = X // 4, Y // 4
+        scale = min(cx // 5, cy // 5)
+        # Resize so smaller dim is 5
 
-        sumframe = np.sum(dframe, axis=2)
-        # TODO: Only search in corner area
-        widx = max_in_orig((np.std(dframe, axis=2) < 35) * sumframe)
-        ridx = max_in_orig(2 * dframe[:, :, 0] - sumframe)
-        gidx = max_in_orig(2 * dframe[:, :, 1] - sumframe)
-        bidx = max_in_orig(2 * dframe[:, :, 2] - sumframe)
+        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)
 
-        # Flood fill corners
-        def flood_fill(s):
-            # TODO: make this faster
-            vis = np.full((X, Y), False)
-            vis[s] = True
-            queue = collections.deque([s])
-            pos = np.array(s)
-            col = np.copy(raw_frame[s])
-            n = 1
-            while len(queue) > 0:
-                u = queue.popleft()
-                for d in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
-                    v = (u[0] + d[0], u[1] + d[1])
-                    if (
-                        0 <= v[0] < X
-                        and 0 <= v[1] < Y
-                        and not vis[v]
-                        and np.linalg.norm(raw_frame[v] - raw_frame[s]) < 125
-                    ):
-                        vis[v] = True
-                        pos += np.array(v)
-                        col += raw_frame[v]
-                        n += 1
-                        queue.append(v)
-            # plt.imshow(raw_frame.astype(np.uint8))
-            # plt.scatter(*reversed(np.where(vis)))
-            # plt.scatter(pos[1] / n, pos[0] / n)
-            # plt.show()
-            return pos / n, col / n
-
-        widx, wcol = flood_fill(widx)
-        ridx, rcol = flood_fill(ridx)
-        gidx, gcol = flood_fill(gidx)
-        bidx, bcol = flood_fill(bidx)
+        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
+        gidx, gcol = find_corner(raw_frame[X - cx :, :cy], lambda B: B[:, :, 1] - B[:, :, 2] - B[:, :, 0])
+        gidx[0] += X - cx
+        bidx, bcol = find_corner(raw_frame[X - cx :, Y - cy :], lambda B: B[:, :, 2] - B[:, :, 0] - B[:, :, 1])
+        bidx[0] += X - cx
+        bidx[1] += Y - cy
 
         # Find basis of color space
         origin = (rcol + gcol + bcol - wcol) / 2
@@ -134,23 +102,28 @@ while data is None:
         # plt.imshow(raw_frame.astype(np.uint8))
         # plt.show()
 
-        raw_color_frame = raw_frame[np.round(xp).astype(np.int64), np.round(yp).astype(np.int64), :]
-        # color_frame = raw_color_frame.astype(np.uint8)
-        color_frame = np.clip(np.squeeze(F @ (raw_color_frame - origin)[..., np.newaxis]), 0, 255).astype(np.uint8)
-        frame = (
-            (color_frame[:, :, 0] >> 5) + (color_frame[:, :, 1] >> 2 & 0b00111000) + (color_frame[:, :, 2] & 0b11000000)
-        )
-        frame_data = np.concatenate(
+        frame = raw_frame[
+            np.clip(np.round(xp).astype(np.int64), 0, X - 1), np.clip(np.round(yp).astype(np.int64), 0, Y - 1), :
+        ]
+        frame = np.clip(np.squeeze(F @ (frame - origin)[..., np.newaxis]), 0, 255).astype(np.uint8)
+        frame = (frame[:, :, 0] >> 7) + (frame[:, :, 1] >> 5 & 0b0110) + (frame[:, :, 2] >> 4 & 0b1000)
+        frame = np.concatenate(
             (
                 frame[:cheight, cwidth : args.width - cwidth].flatten(),
                 frame[cheight : args.height - cheight].flatten(),
                 frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
             )
         )
-        data = decoder.decode(bytes(rsc.decode(frame_data ^ frame_xor)[0]))
+        frame = ((frame[::2] << 4) + frame[1::2]) ^ frame_xor
+        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]))
         print("Decoded frame")
-    except Exception as e:
-        print(e)
+    except KeyboardInterrupt:
+        sys.exit()
+    except:
+        traceback.print_exc()
 with open(args.output, "wb") as f:
     f.write(data)
 cap.release()
diff --git a/encoder.py b/encoder.py
index 4b71bf4..0c68866 100644
--- a/encoder.py
+++ b/encoder.py
@@ -23,11 +23,11 @@ args = parser.parse_args()
 cheight = cwidth = max(args.height // 10, args.width // 10)
 midwidth = args.width - 2 * cwidth
 frame_size = args.height * args.width - 4 * cheight * cwidth
+# Divide by 2 for 4-bit color
 frame_xor = np.arange(frame_size // 2, dtype=np.uint8)
 # reedsolo breaks message into 255-byte chunks
 # raptorq can add up to 4 extra bytes
-# Divide by 2 for 4-bit color
-rs_size = frame_size - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
+rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
 
 with open(args.input, "rb") as f:
     data = f.read()
@@ -55,7 +55,7 @@ def get_frame():
     # 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(np.reshape(frame, (len(frame) // 255, 255)), "F")[: frame_size // 2] ^ frame_xor
+    frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2] ^ frame_xor
     frame = np.ravel(np.column_stack((frame >> 4, frame & 0b1111)))
     frame = np.concatenate(
         (