diff options
Diffstat (limited to 'encoder.py')
| -rw-r--r-- | encoder.py | 39 |
1 files changed, 21 insertions, 18 deletions
@@ -23,22 +23,24 @@ 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 -frame_xor = np.arange(frame_size, dtype=np.uint8) +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 -rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4 +# Divide by 2 for 4-bit color +rs_size = frame_size - (frame_size // 2 + 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) +# This formula is wrong 🤷 packets = encoder.get_encoded_packets(int(len(data) / rs_size * args.level)) # Make corners -wcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b11111111), ((0, 1), (0, 1))) -rcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b00000111), ((0, 1), (1, 0))) -gcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b00111000), ((1, 0), (0, 1))) -bcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b11000000), ((1, 0), (1, 0))) +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))) print("Data length:", len(data)) print("Packets:", len(packets)) @@ -48,29 +50,30 @@ idx = 0 def get_frame(): global idx - frame_data = np.array(rsc.encode(packets[idx])) - # Pad frame to fit frame_size since raptorq might not add 4 bytes - frame_data = np.pad(frame_data, (0, frame_size - len(frame_data))) ^ frame_xor - if idx == 0: - print(list(frame_data)) + frame = np.array(rsc.encode(packets[idx])) idx = (idx + 1) % len(packets) + # 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(np.column_stack((frame >> 4, frame & 0b1111))) frame = np.concatenate( ( np.concatenate( - (wcorner, frame_data[: cheight * midwidth].reshape((cheight, midwidth)), rcorner), + (wcorner, frame[: cheight * midwidth].reshape((cheight, midwidth)), rcorner), axis=1, ), - frame_data[cheight * midwidth : frame_size - cheight * midwidth].reshape( + frame[cheight * midwidth : frame_size - cheight * midwidth].reshape( (args.height - 2 * cheight, args.width) ), np.concatenate( - (gcorner, frame_data[frame_size - cheight * midwidth :].reshape((cheight, midwidth)), bcorner), + (gcorner, frame[frame_size - cheight * midwidth :].reshape((cheight, midwidth)), bcorner), axis=1, ), ) ) return np.stack( - ((frame & 0b00000111) * 255 // 7, (frame >> 3 & 0b00000111) * 255 // 7, (frame >> 6) * 255 // 3), + ((frame & 0b0001) * 255, (frame >> 1 & 0b0011) * 255 / 3, (frame >> 3) * 255), axis=-1, ).astype(np.uint8) @@ -102,7 +105,7 @@ if args.output: width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) hscale = height // args.height wscale = width // args.width - out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"mp4v"), args.fps, (width, height)) + out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"RGBA"), args.fps, (width, height)) while cap.isOpened(): ret, frame = cap.read() if not ret: @@ -119,9 +122,9 @@ if args.output: ) ) else: - out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"mp4v"), args.fps, (4* args.width, 4*args.height)) + out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"RGBA"), args.fps, (args.width, args.height)) for _ in packets: - out.write(cv2.cvtColor(np.repeat(np.repeat(get_frame(), 4, 0), 4, 1), cv2.COLOR_RGB2BGR)) + out.write(cv2.cvtColor(get_frame(), cv2.COLOR_RGB2BGR)) else: input("Seizure warning!") app = QApplication([]) |