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