3-bit color

2 files changed, 35 insertions, 40 deletions

diff --git a/decoder.py b/decoder.py
index 2b09081..efa9a81 100644
--- a/decoder.py
+++ b/decoder.py
@@ -20,11 +20,12 @@ 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 // 2, dtype=np.uint8)
-rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
+frame_bytes = frame_size * 3 // 8
+frame_xor = np.arange(frame_bytes, dtype=np.uint8)
+rs_bytes = frame_bytes - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4
 
 rsc = RSCodec(int(args.level * 255))
-decoder = Decoder.with_defaults(args.size, rs_size)
+decoder = Decoder.with_defaults(args.size, rs_bytes)
 
 
 def find_corner(A, f):
@@ -38,8 +39,8 @@ def find_corner(A, f):
         np.empty(0),
         tuple(reversed(guess)),
         0,
-        (50, 50, 50),
-        (50, 50, 50),
+        (100, 100, 100),
+        (100, 100, 100),
         cv2.FLOODFILL_MASK_ONLY + cv2.FLOODFILL_FIXED_RANGE,
     )[2][1:-1, 1:-1].astype(bool)
     return np.average(np.where(mask), axis=1), np.average(A[mask], axis=0).astype(np.float64)
@@ -113,8 +114,7 @@ while data is None:
         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.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 128).astype(np.uint8)
         frame = np.concatenate(
             (
                 frame[:cheight, cwidth : args.width - cwidth].flatten(),
@@ -122,11 +122,8 @@ while data is None:
                 frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
             )
         )
-        frame = (frame[::2] << 4) + frame[1::2]
-        # 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]
-        erase_pos = bytearray(np.where(frame == 0)[0]) if args.erasure else bytearray()
-        data = decoder.decode(bytes(rsc.decode(bytearray(frame ^ frame_xor), erase_pos=erase_pos)[0]))
+        # erase_pos = bytearray(np.where(frame == 0)[0]) if args.erasure else bytearray()
+        data = decoder.decode(bytes(rsc.decode(bytearray(np.packbits(frame) ^ frame_xor))[0]))
         print("Decoded frame")
     except KeyboardInterrupt:
         sys.exit()
diff --git a/encoder.py b/encoder.py
index 6b7b3d5..914e396 100644
--- a/encoder.py
+++ b/encoder.py
@@ -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))