diff options
Diffstat (limited to 'decoder.py')
| -rw-r--r-- | decoder.py | 145 |
1 files changed, 81 insertions, 64 deletions
@@ -1,4 +1,6 @@ import argparse +from multiprocessing.pool import Pool +from threading import Thread import time import cv2 import numpy as np @@ -23,6 +25,9 @@ rs_bytes = frame_bytes - (frame_bytes + 254) // 255 * int(args.level * 255) - 4 rsc = RSCodec(int(args.level * 255)) decoder = Decoder.with_defaults(args.size, rs_bytes) +data = None +decoded = 0 +start_time = 0 def find_corner(A, f): @@ -43,78 +48,90 @@ def find_corner(A, f): return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64) -if isinstance(args.input, str) and args.input.isdecimal(): - args.input = int(args.input) -cap = cv2.VideoCapture(args.input) -data = None -start_time = 0 -decoded = 0 -while data is None: - try: - ret, raw_frame = cap.read() - if not ret: - print("End of stream") - break - # raw_frame is a uint8 BE CAREFUL - cv2.imshow("", raw_frame) - cv2.waitKey(1) - raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB) - - # Find positions and colors of corners - X, Y = raw_frame.shape[:2] - cx, cy = X // 4, Y // 4 - widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: np.sum(B, axis=2) - 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]) - 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 +def rsc_decode(x): + return rsc.decode(x)[0] - # Find basis of color space - origin = (rcol + gcol + bcol - wcol) / 2 - rcol -= origin - gcol -= origin - bcol -= origin - F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T) - cch = cheight / 2 - 1 - ccw = cwidth / 2 - 1 - M = cv2.getPerspectiveTransform( - np.float32([np.flip(widx), np.flip(ridx), np.flip(gidx), np.flip(bidx)]), - np.float32( - [ - [ccw, cch], - [args.width - ccw - 1, cch], - [ccw, args.height - cch - 1], - [args.width - ccw - 1, args.height - cch - 1], - ] - ), - ) - frame = cv2.warpPerspective(raw_frame, M, (args.width, args.height)) - # Convert to new color space - frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 160).astype(np.uint8) - frame = np.packbits( - np.concatenate( - ( - frame[:cheight, cwidth : args.width - cwidth].flatten(), - frame[cheight : args.height - cheight].flatten(), - frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(), - ) - ) - )[:frame_bytes] - reshape_len = frame_bytes // 255 * 255 - frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(255, reshape_len // 255), "F") - data = decoder.decode(bytes(rsc.decode(bytearray(frame ^ frame_xor))[0][: args.psize])) +def ecc_decode(frame): + global data + global decoded + global start_time + try: + rsc_frame = bytearray(frame ^ frame_xor) + with Pool(processes=1) as p: + rsc_data = p.apply(rsc_decode, (rsc_frame,)) + data = decoder.decode(bytes(rsc_data[: args.psize])) decoded += 1 if start_time == 0: start_time = time.time() print("Decoded frame") - except KeyboardInterrupt: - break except Exception as e: print(e) + + +if isinstance(args.input, str) and args.input.isdecimal(): + args.input = int(args.input) +cap = cv2.VideoCapture(args.input) +while data is None: + frame_start_time = time.time() + ret, raw_frame = cap.read() + if not ret: + print("End of stream") + break + # raw_frame is a uint8 BE CAREFUL + cv2.imshow("", raw_frame) + cv2.waitKey(1) + raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB) + + # Find positions and colors of corners + X, Y = raw_frame.shape[:2] + cx, cy = X // 4, Y // 4 + widx, wcol = find_corner(raw_frame[:cx, :cy], lambda B: np.sum(B, axis=2) - 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]) + 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 + rcol -= origin + gcol -= origin + bcol -= origin + F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T) + + cch = cheight / 2 - 1 + ccw = cwidth / 2 - 1 + M = cv2.getPerspectiveTransform( + np.float32([np.flip(widx), np.flip(ridx), np.flip(gidx), np.flip(bidx)]), + np.float32( + [ + [ccw, cch], + [args.width - ccw - 1, cch], + [ccw, args.height - cch - 1], + [args.width - ccw - 1, args.height - cch - 1], + ] + ), + ) + frame = cv2.warpPerspective(raw_frame, M, (args.width, args.height)) + # Convert to new color space + frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 160).astype(np.uint8) + frame = np.packbits( + np.concatenate( + ( + frame[:cheight, cwidth : args.width - cwidth].flatten(), + frame[cheight : args.height - cheight].flatten(), + frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(), + ) + ) + )[:frame_bytes] + reshape_len = frame_bytes // 255 * 255 + frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(255, reshape_len // 255), "F") + # Run ECC decoding in separate thread + Thread(target=ecc_decode, args=(frame,)).start() + print(time.time() - frame_start_time) cap.release() print(decoded) with open(args.output, "wb") as f: |