diff options
Diffstat (limited to 'decoder.py')
| -rw-r--r-- | decoder.py | 34 |
1 files changed, 17 insertions, 17 deletions
@@ -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() |