aboutsummaryrefslogtreecommitdiff
path: root/decoder.py
diff options
context:
space:
mode:
authorAnthony Wang2024-04-23 13:23:56 -0400
committerAnthony Wang2024-04-23 13:23:56 -0400
commit4b698588160ca5a9e8dacf8de105f50b0f7e3c6c (patch)
tree1a12562e2004140ea451bc5d31203f341c44919f /decoder.py
parent3d940c0a1c9b5eee2adbbdd6f0a358e782a4df06 (diff)
Move find_corner to global namespace, use erase_pos for better ECC
Diffstat (limited to 'decoder.py')
-rw-r--r--decoder.py34
1 files changed, 17 insertions, 17 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()