4-bit color for encoder, fix a ton of bugs
This commit is contained in:
parent
dee89fcb30
commit
3d940c0a1c
SSH key fingerprint:
SHA256:B5ADfMCqd2M7d/jtXDoihAV/yfXOAbWWri9+GdCN4hQ
|
|
@ -3,8 +3,5 @@
|
|||
Formatting: `black -l 120 *.py`
|
||||
|
||||
TODO:
|
||||
- Only search corner for corner
|
||||
- Fast BFS using numpy (OpenCV floodfill?)
|
||||
- Limit BFS time
|
||||
- Spread out RS bytes (still need to impl in decoder)
|
||||
- 4bit color (still need to impl in decoder)
|
||||
- Write better ECC
|
||||
|
|
|
|||
107
decoder.py
107
decoder.py
|
|
@ -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()
|
||||
cx, cy = X // 4, Y // 4
|
||||
scale = min(cx // 5, cy // 5)
|
||||
# Resize so smaller dim is 5
|
||||
|
||||
def max_in_orig(x):
|
||||
return tuple(np.array(np.unravel_index(np.argmax(x), x.shape)) * scale + scale // 2)
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
# 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()
|
||||
|
|
|
|||
|
|
@ -23,11 +23,11 @@ args = parser.parse_args()
|
|||
cheight = cwidth = max(args.height // 10, args.width // 10)
|
||||
midwidth = args.width - 2 * cwidth
|
||||
frame_size = args.height * args.width - 4 * cheight * cwidth
|
||||
# Divide by 2 for 4-bit color
|
||||
frame_xor = np.arange(frame_size // 2, dtype=np.uint8)
|
||||
# reedsolo breaks message into 255-byte chunks
|
||||
# raptorq can add up to 4 extra bytes
|
||||
# Divide by 2 for 4-bit color
|
||||
rs_size = frame_size - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
|
||||
rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
|
||||
|
||||
with open(args.input, "rb") as f:
|
||||
data = f.read()
|
||||
|
|
@ -55,7 +55,7 @@ def get_frame():
|
|||
# Add 4 bytes, pad frame to be multiple of 255
|
||||
frame = np.pad(frame, (0, (len(frame) + 258) // 255 * 255 - len(frame)))
|
||||
# Space out elements in each size 255 chunk
|
||||
frame = np.ravel(np.reshape(frame, (len(frame) // 255, 255)), "F")[: frame_size // 2] ^ frame_xor
|
||||
frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2] ^ frame_xor
|
||||
frame = np.ravel(np.column_stack((frame >> 4, frame & 0b1111)))
|
||||
frame = np.concatenate(
|
||||
(
|
||||
|
|
|
|||
Loading…
Reference in a new issue