Mix frames with all vidframe pixels that are 0 or 255

3 files changed, 17 insertions, 19 deletions

diff --git a/README.md b/README.md
index d465263..98441c9 100644
--- a/README.md
+++ b/README.md
@@ -12,10 +12,14 @@ pip install --upgrade reedsolo --no-binary "reedsolo" --no-cache --config-settin
 
 Encode: `python encoder.py -i in`
 
-Play video: `mpv --scale=nearest --fullscreen --loop --no-keepaspect vid.mkv`
+Play video (SEIZURE WARNING): `mpv --scale=nearest --fullscreen --loop --no-keepaspect vid.mkv`
 
 Copy the flags printed by the encoder and pass them to the decoder: `python decoder.py FLAGS`
 
 Formatting: `black -l 120 *.py`
 
 Use phone as webcam for higher quality video: `scrcpy --v4l2-sink=/dev/video4 --video-source=camera --no-video-playback --camera-size 1920x1440`
+
+## TODO
+
+- Reorder bits
diff --git a/decoder.py b/decoder.py
index 2fb0183..b0d06ef 100644
--- a/decoder.py
+++ b/decoder.py
@@ -11,11 +11,9 @@ parser.add_argument("-i", "--input", help="camera device index or input video fi
 parser.add_argument("-o", "--output", help="output file for decoded data", default="out")
 parser.add_argument("-x", "--height", help="grid height", default=100, type=int)
 parser.add_argument("-y", "--width", help="grid width", default=100, type=int)
-parser.add_argument("-f", "--fps", help="frame rate", default=30, type=int)
 parser.add_argument("-l", "--level", help="error correction level", default=0.1, type=float)
 parser.add_argument("-s", "--size", help="number of bytes to decode", type=int)
 parser.add_argument("-p", "--psize", help="packet size", type=int)
-parser.add_argument("-e", "--erasure", help="detect erasures", action="store_true")
 args = parser.parse_args()
 
 cheight = cwidth = max(args.height // 10, args.width // 10)
@@ -99,7 +97,6 @@ while data is None:
                 frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
             )
         )
-        # 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][: args.psize]))
         print("Decoded frame")
     except KeyboardInterrupt:
diff --git a/encoder.py b/encoder.py
index d80beae..bd44876 100644
--- a/encoder.py
+++ b/encoder.py
@@ -9,8 +9,8 @@ parser.add_argument("-i", "--input", help="input file")
 parser.add_argument("-o", "--output", help="output video file", default="vid.mkv")
 parser.add_argument("-x", "--height", help="grid height", default=100, type=int)
 parser.add_argument("-y", "--width", help="grid width", default=100, type=int)
-parser.add_argument("-f", "--fps", help="frame rate", default=30, type=int)
 parser.add_argument("-l", "--level", help="error correction level", default=0.1, type=float)
+parser.add_argument("-f", "--fps", help="frame rate", default=30, type=int)
 parser.add_argument("-m", "--mix", help="mix frames with original video", action="store_true")
 args = parser.parse_args()
 
@@ -42,7 +42,7 @@ bcorner = np.pad(np.dstack((zeros, zeros, ones)), ((1, 0), (1, 0), (0, 0)))
 print(f"-x {args.height} -y {args.width} -l {args.level} -s {len(data)} -p {len(packets[0])}", end="")
 
 
-def frame(packet):
+def mkframe(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))
@@ -74,22 +74,19 @@ if args.mix:
     out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (width, height))
     i = 0
     while cap.isOpened():
-        ret, frame = cap.read()
+        ret, vidframe = cap.read()
         if not ret:
             break
-        frame = frame.astype(np.float64) / 255
-        frame[: hscale * cheight, : wscale * cwidth] = 1
-        frame[: hscale * cheight, wscale * (args.width - cwidth) :] = 1
-        frame[hscale * (args.height - cheight) :, : wscale * cwidth] = 1
-        frame[hscale * (args.height - cheight) :, wscale * (args.width - cwidth) :] = 1
-        out.write(
-            cv2.cvtColor(
-                (frame * np.repeat(np.repeat(frame(packets[i]), hscale, 0), wscale, 1)).astype(np.uint8),
-                cv2.COLOR_RGB2BGR,
-            )
-        )
+        vidframe[: hscale * cheight, : wscale * cwidth] = 0
+        vidframe[: hscale * cheight, wscale * (args.width - cwidth) :] = 0
+        vidframe[hscale * (args.height - cheight) :, : wscale * cwidth] = 0
+        vidframe[hscale * (args.height - cheight) :, wscale * (args.width - cwidth) :] = 0
+
+        frame = np.repeat(np.repeat(mkframe(packets[i]), hscale, 0), wscale, 1)
+        frame[vidframe % 255 != 0] = 0
+        out.write(cv2.cvtColor(frame, 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 packet in packets:
-        out.write(cv2.cvtColor(frame(packet), cv2.COLOR_RGB2BGR))
+        out.write(cv2.cvtColor(mkframe(packet), cv2.COLOR_RGB2BGR))