diff options
Diffstat (limited to 'encoder.py')
| -rw-r--r-- | encoder.py | 54 |
1 files changed, 26 insertions, 28 deletions
@@ -18,56 +18,52 @@ 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) +# Divide by 8 / 3 for 3-bit color +frame_bytes = frame_size * 3 // 8 +frame_xor = np.arange(frame_bytes, dtype=np.uint8) # reedsolo breaks message into 255-byte chunks # raptorq can add up to 4 extra bytes -rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4 +rs_bytes = frame_bytes - (frame_bytes + 254) // 255 * int(args.level * 255) - 4 with open(args.input, "rb") as f: data = f.read() rsc = RSCodec(int(args.level * 255)) -encoder = Encoder.with_defaults(data, rs_size) -packets = encoder.get_encoded_packets(int(len(data) / rs_size * (1 / (1 - args.level) - 1))) +encoder = Encoder.with_defaults(data, rs_bytes) +packets = encoder.get_encoded_packets(int(len(data) / rs_bytes * (1 / (1 - args.level) - 1))) # Make corners -wcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b1111), ((0, 1), (0, 1))) -rcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b0001), ((0, 1), (1, 0))) -gcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b0110), ((1, 0), (0, 1))) -bcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b1000), ((1, 0), (1, 0))) +ones = np.ones((cheight - 1, cwidth - 1)) +zeros = np.zeros((cheight - 1, cwidth - 1)) +wcorner = np.pad(np.dstack((ones, ones, ones)), ((0, 1), (0, 1), (0, 0))) +rcorner = np.pad(np.dstack((ones, zeros, zeros)), ((0, 1), (1, 0), (0, 0))) +gcorner = np.pad(np.dstack((zeros, ones, zeros)), ((1, 0), (0, 1), (0, 0))) +bcorner = np.pad(np.dstack((zeros, zeros, ones)), ((1, 0), (1, 0), (0, 0))) print("Data length:", len(data)) print("Packets:", len(packets)) -idx = 0 - -def get_frame(): - global idx - frame = np.array(rsc.encode(bytearray(packets[idx]))) - idx = (idx + 1) % len(packets) - frame = np.pad(frame, (0, frame_size // 2 - len(frame))) ^ frame_xor - # Pad frame to be multiple of 255 - # frame = np.pad(frame, (0, (len(frame) + 254) // 255 * 255 - len(frame))) - # Space out elements in each size 255 chunk - # frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2] - frame = np.ravel(np.column_stack((frame >> 4, frame & 0b1111))) +def frame(packet): + frame = np.array(rsc.encode(bytearray(packet))) + frame = np.unpackbits(np.pad(frame, (0, frame_bytes - len(frame))) ^ frame_xor) + frame = np.pad(frame, (0, (3 - len(frame)) % 3)) + frame = np.reshape(frame, (frame_size, 3)) frame = np.concatenate( ( np.concatenate( - (wcorner, frame[: cheight * midwidth].reshape((cheight, midwidth)), rcorner), + (wcorner, frame[: cheight * midwidth].reshape((cheight, midwidth, 3)), rcorner), axis=1, ), frame[cheight * midwidth : frame_size - cheight * midwidth].reshape( - (args.height - 2 * cheight, args.width) + (args.height - 2 * cheight, args.width, 3) ), np.concatenate( - (gcorner, frame[frame_size - cheight * midwidth :].reshape((cheight, midwidth)), bcorner), + (gcorner, frame[frame_size - cheight * midwidth :].reshape((cheight, midwidth, 3)), bcorner), axis=1, ), ) ) - return np.stack(((frame & 0b0001) * 255, (frame >> 1 & 0b0011) * 85, (frame >> 3) * 255), axis=-1).astype(np.uint8) + return frame.astype(np.uint8) * 255 if args.mix: @@ -77,6 +73,7 @@ if args.mix: hscale = height // args.height wscale = width // args.width out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (width, height)) + i = 0 while cap.isOpened(): ret, frame = cap.read() if not ret: @@ -88,11 +85,12 @@ if args.mix: frame[hscale * (args.height - cheight) :, wscale * (args.width - cwidth) :] = 1 out.write( cv2.cvtColor( - (frame * np.repeat(np.repeat(get_frame(), hscale, 0), wscale, 1)).astype(np.uint8), + (frame * np.repeat(np.repeat(frame(packets[i]), hscale, 0), wscale, 1)).astype(np.uint8), cv2.COLOR_RGB2BGR, ) ) + i = (i + 1) % len(packets) else: out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (args.width, args.height)) - for _ in packets: - out.write(cv2.cvtColor(get_frame(), cv2.COLOR_RGB2BGR)) + for packet in packets: + out.write(cv2.cvtColor(frame(packet), cv2.COLOR_RGB2BGR)) |