Mix frames with all vidframe pixels that are 0 or 255

1 files changed, 12 insertions, 15 deletions

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))