Make corner finding more robust

1 files changed, 19 insertions, 12 deletions

diff --git a/decoder.py b/decoder.py
index dbe9751..2b09081 100644
--- a/decoder.py
+++ b/decoder.py
@@ -2,6 +2,7 @@ import argparse
 import sys
 import traceback
 import cv2
+import matplotlib.pyplot as plt
 import numpy as np
 from creedsolo import RSCodec
 from raptorq import Decoder
@@ -28,13 +29,19 @@ 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)
+    # Resize so smaller dim is 8
+    scale = min(cx // 8, cy // 8)
     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)
+    guess = np.array(np.unravel_index(np.argmax(f(B.astype(np.float64))), B.shape[:2])) * scale + scale // 2
+    mask = cv2.floodFill(
+        A,
+        np.empty(0),
+        tuple(reversed(guess)),
+        0,
+        (50, 50, 50),
+        (50, 50, 50),
+        cv2.FLOODFILL_MASK_ONLY + cv2.FLOODFILL_FIXED_RANGE,
+    )[2][1:-1, 1:-1].astype(bool)
     return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64)
 
 
@@ -54,8 +61,8 @@ while data is None:
         raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB)
 
         X, Y = raw_frame.shape[:2]
-        cx, cy = X // 4, Y // 4
-        widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: (np.std(B, axis=2) < 35) * np.sum(B, axis=2))
+        cx, cy = X // 3, Y // 3
+        widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: np.sum(B, axis=2) - np.std(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])
@@ -116,10 +123,10 @@ while data is None:
             )
         )
         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]
-        erase_pos = list(np.where(frame == 0)[0]) if args.erasure else []
-        data = decoder.decode(bytes(rsc.decode(frame ^ frame_xor, erase_pos=erase_pos)[0]))
+        # 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]
+        erase_pos = bytearray(np.where(frame == 0)[0]) if args.erasure else bytearray()
+        data = decoder.decode(bytes(rsc.decode(bytearray(frame ^ frame_xor), erase_pos=erase_pos)[0]))
         print("Decoded frame")
     except KeyboardInterrupt:
         sys.exit()