Short args, clean up code

2 files changed, 164 insertions, 192 deletions

diff --git a/decoder.py b/decoder.py
index 649220c..f9829d5 100644
--- a/decoder.py
+++ b/decoder.py
@@ -1,5 +1,6 @@
 import argparse
 import collections
+import sys
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
@@ -7,13 +8,13 @@ from creedsolo import RSCodec
 from raptorq import Decoder
 
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-parser.add_argument("file", help="output file for decoded data")
-parser.add_argument("--len", help="number of bytes to decode", default=2**16, type=int)
-parser.add_argument("--height", help="grid height", default=100, type=int)
-parser.add_argument("--width", help="grid width", default=100, type=int)
-parser.add_argument("--fps", help="framerate", default=30, type=int)
-parser.add_argument("--level", help="error correction level", default=0.1, type=float)
-parser.add_argument("--device", help="camera device index", default=1, type=int)
+parser.add_argument("-i", "--input", help="camera device index or input video file", default=0)
+parser.add_argument("-o", "--output", help="output file for decoded data")
+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", default=2**16, type=int)
 args = parser.parse_args()
 
 cheight = cwidth = max(args.height // 10, args.width // 10)
@@ -22,143 +23,127 @@ frame_xor = np.arange(frame_size, dtype=np.uint8)
 rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4
 
 rsc = RSCodec(int(args.level * 255))
-decoder = Decoder.with_defaults(args.len, rs_size)
+decoder = Decoder.with_defaults(args.size, rs_size)
+
+if args.input.isdecimal():
+    args.input = int(args.input)
+cap = cv2.VideoCapture(args.input)
 data = None
-# cap = cv2.VideoCapture(args.device)
 while data is None:
-    # ret, raw_frame = cap.read()
-    # if not ret:
-    #    continue
-    # TODO: Try decoding saved videos
-    raw_frame = cv2.cvtColor(
-        cv2.imread("/home/a/Pictures/Camera/IMG_20240422_000849_027.jpg"),
-        cv2.COLOR_BGR2RGB,
-    ).astype(np.float64)
-
-    X, Y = raw_frame.shape[:2]
-    scale = min(X // 20, Y // 20)
-    # Resize so smaller dim is 20
-    # Use fast default interpolation for factor of 4
-    # Then switch to good slow interpolation
-    dframe = cv2.resize(
-        cv2.resize(raw_frame, (Y // 4, X // 4)),
-        (Y // scale, X // scale),  # OpenCV swaps them
-        interpolation=cv2.INTER_AREA,
-    )
-    plt.imshow(dframe.astype(np.uint8))
-    plt.show()
+    try:
+        ret, raw_frame = cap.read()
+        if not ret:
+            print("End of stream")
+            sys.exit(1)
+        raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB).astype(np.float64)
 
-    def max_in_orig(x):
-        return tuple(
-            np.array(np.unravel_index(np.argmax(x), x.shape)) * scale + scale // 2
+        X, Y = raw_frame.shape[:2]
+        scale = min(X // 20, Y // 20)
+        # Resize so smaller dim is 20
+        # Use fast default interpolation for factor of 4
+        # Then switch to good slow interpolation
+        dframe = cv2.resize(
+            cv2.resize(raw_frame, (Y // 4, X // 4)),
+            (Y // scale, X // scale),  # OpenCV swaps them
+            interpolation=cv2.INTER_AREA,
         )
+        # plt.imshow(dframe.astype(np.uint8))
+        # plt.show()
 
-    sumframe = np.sum(dframe, axis=2)
-    # TODO: Only search in corner area
-    widx = max_in_orig((np.std(dframe, axis=2) < 35) * sumframe)
-    ridx = max_in_orig(2 * dframe[:, :, 0] - sumframe)
-    gidx = max_in_orig(2 * dframe[:, :, 1] - sumframe)
-    bidx = max_in_orig(2 * dframe[:, :, 2] - sumframe)
+        def max_in_orig(x):
+            return tuple(np.array(np.unravel_index(np.argmax(x), x.shape)) * scale + scale // 2)
 
-    # Flood fill corners
-    def flood_fill(s):
-        vis = np.full((X, Y), False)
-        vis[s] = True
-        queue = collections.deque([s])
-        pos = np.array(s)
-        col = np.copy(raw_frame[s])
-        n = 1
-        while len(queue) > 0:
-            u = queue.popleft()
-            for d in [(5, 0), (0, 5), (-5, 0), (0, -5)]:
-                v = (u[0] + d[0], u[1] + d[1])
-                if (
-                    0 <= v[0] < X
-                    and 0 <= v[1] < Y
-                    and not vis[v]
-                    and np.linalg.norm(raw_frame[v] - raw_frame[s]) < 100
-                ):
-                    vis[v] = True
-                    pos += np.array(v)
-                    col += raw_frame[v]
-                    n += 1
-                    queue.append(v)
-        plt.imshow(raw_frame.astype(np.uint8))
-        plt.scatter(*reversed(np.where(vis)))
-        plt.scatter(pos[1] / n, pos[0] / n)
-        plt.show()
-        return pos / n, col / n
+        sumframe = np.sum(dframe, axis=2)
+        # TODO: Only search in corner area
+        widx = max_in_orig((np.std(dframe, axis=2) < 35) * sumframe)
+        ridx = max_in_orig(2 * dframe[:, :, 0] - sumframe)
+        gidx = max_in_orig(2 * dframe[:, :, 1] - sumframe)
+        bidx = max_in_orig(2 * dframe[:, :, 2] - sumframe)
 
-    widx, wcol = flood_fill(widx)
-    ridx, rcol = flood_fill(ridx)
-    gidx, gcol = flood_fill(gidx)
-    bidx, bcol = flood_fill(bidx)
+        # Flood fill corners
+        def flood_fill(s):
+            # TODO: make this faster
+            vis = np.full((X, Y), False)
+            vis[s] = True
+            queue = collections.deque([s])
+            pos = np.array(s)
+            col = np.copy(raw_frame[s])
+            n = 1
+            while len(queue) > 0:
+                u = queue.popleft()
+                for d in [(1, 0), (0, 1), (-1, 0), (0, -1)]:
+                    v = (u[0] + d[0], u[1] + d[1])
+                    if 0 <= v[0] < X and 0 <= v[1] < Y and not vis[v] and np.linalg.norm(raw_frame[v] - raw_frame[s]) < 125:
+                        vis[v] = True
+                        pos += np.array(v)
+                        col += raw_frame[v]
+                        n += 1
+                        queue.append(v)
+            # plt.imshow(raw_frame.astype(np.uint8))
+            # plt.scatter(*reversed(np.where(vis)))
+            # plt.scatter(pos[1] / n, pos[0] / n)
+            # plt.show()
+            return pos / n, col / n
 
-    # Find basis of color space
-    origin = (rcol + gcol + bcol - wcol) / 2
-    rcol -= origin
-    gcol -= origin
-    bcol -= origin
-    print(origin, rcol, gcol, bcol)
-    F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T)
+        widx, wcol = flood_fill(widx)
+        ridx, rcol = flood_fill(ridx)
+        gidx, gcol = flood_fill(gidx)
+        bidx, bcol = flood_fill(bidx)
 
-    # Dumb perspective transform
-    xv = np.linspace(
-        -(cheight / 2 - 1) / (args.height - cheight + 1),
-        1 + (cheight / 2 - 1) / (args.height - cheight + 1),
-        args.height,
-    )
-    yv = np.linspace(
-        -(cwidth / 2 - 1) / (args.width - cwidth + 1),
-        1 + (cwidth / 2 - 1) / (args.width - cwidth + 1),
-        args.width,
-    )
-    xp = (
-        np.outer(1 - xv, 1 - yv) * widx[0]
-        + np.outer(1 - xv, yv) * ridx[0]
-        + np.outer(xv, 1 - yv) * gidx[0]
-        + np.outer(xv, yv) * bidx[0]
-    )
-    yp = (
-        np.outer(1 - xv, 1 - yv) * widx[1]
-        + np.outer(1 - xv, yv) * ridx[1]
-        + np.outer(xv, 1 - yv) * gidx[1]
-        + np.outer(xv, yv) * bidx[1]
-    )
+        # Find basis of color space
+        origin = (rcol + gcol + bcol - wcol) / 2
+        rcol -= origin
+        gcol -= origin
+        bcol -= origin
+        F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T)
+
+        # Dumb perspective transform
+        xv = np.linspace(
+            -(cheight / 2 - 1) / (args.height - cheight + 1),
+            1 + (cheight / 2 - 1) / (args.height - cheight + 1),
+            args.height,
+        )
+        yv = np.linspace(
+            -(cwidth / 2 - 1) / (args.width - cwidth + 1),
+            1 + (cwidth / 2 - 1) / (args.width - cwidth + 1),
+            args.width,
+        )
+        xp = (
+            np.outer(1 - xv, 1 - yv) * widx[0]
+            + np.outer(1 - xv, yv) * ridx[0]
+            + np.outer(xv, 1 - yv) * gidx[0]
+            + np.outer(xv, yv) * bidx[0]
+        )
+        yp = (
+            np.outer(1 - xv, 1 - yv) * widx[1]
+            + np.outer(1 - xv, yv) * ridx[1]
+            + np.outer(xv, 1 - yv) * gidx[1]
+            + np.outer(xv, yv) * bidx[1]
+        )
 
-    plt.scatter(widx[1], widx[0])
-    plt.scatter(ridx[1], ridx[0])
-    plt.scatter(gidx[1], gidx[0])
-    plt.scatter(bidx[1], bidx[0])
-    plt.scatter(yp, xp)
-    plt.imshow(raw_frame.astype(np.uint8))
-    plt.show()
-    print(111111111, xp)
-    print(111111111, yp)
+        # plt.scatter(widx[1], widx[0])
+        # plt.scatter(ridx[1], ridx[0])
+        # plt.scatter(gidx[1], gidx[0])
+        # plt.scatter(bidx[1], bidx[0])
+        # plt.scatter(yp, xp)
+        # plt.imshow(raw_frame.astype(np.uint8))
+        # plt.show()
 
-    raw_color_frame = raw_frame[xp.astype(np.int64), yp.astype(np.int64), :]
-    print(raw_color_frame)
-    color_frame = np.clip(
-        np.squeeze(F @ (raw_color_frame - origin)[..., np.newaxis]), 0, 255
-    ).astype(np.uint8)
-    print(color_frame)
-    frame = (
-        (color_frame[:, :, 0] >> 5)
-        + (color_frame[:, :, 1] >> 3 & 0b00111000)
-        + (color_frame[:, :, 2] & 0b11000000)
-    )
-    frame_data = np.concatenate(
-        (
-            frame[:cheight, cwidth : args.width - cwidth].flatten(),
-            frame[cheight : args.height - cheight].flatten(),
-            frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
+        raw_color_frame = raw_frame[np.round(xp).astype(np.int64), np.round(yp).astype(np.int64), :]
+        # color_frame = raw_color_frame.astype(np.uint8)
+        color_frame = np.clip(np.squeeze(F @ (raw_color_frame - origin)[..., np.newaxis]), 0, 255).astype(np.uint8)
+        frame = (color_frame[:, :, 0] >> 5) + (color_frame[:, :, 1] >> 2 & 0b00111000) + (color_frame[:, :, 2] & 0b11000000)
+        frame_data = np.concatenate(
+            (
+                frame[:cheight, cwidth : args.width - cwidth].flatten(),
+                frame[cheight : args.height - cheight].flatten(),
+                frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
+            )
         )
-    )
-    print(list(frame_data))
-    tmp = rsc.decode(frame_data ^ frame_xor)
-    # print(tmp, list(tmp[2]), bytes(tmp[0]))
-    data = decoder.decode(bytes(tmp))
-    break
-with open(args.file, "wb") as f:
+        data = decoder.decode(bytes(rsc.decode(frame_data ^ frame_xor)[0]))
+        print("Decoded frame")
+    except Exception as e:
+        print(e)
+with open(args.output, "wb") as f:
     f.write(data)
 cap.release()
diff --git a/encoder.py b/encoder.py
index b20cd56..5752d55 100644
--- a/encoder.py
+++ b/encoder.py
@@ -10,27 +10,17 @@ from PIL import Image, ImageQt
 from raptorq import Encoder
 
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-parser.add_argument("file", help="input file")
-parser.add_argument("--height", help="grid height", default=100, type=int)
-parser.add_argument("--width", help="grid width", default=100, type=int)
-parser.add_argument("--fps", help="framerate", default=30, type=int)
-parser.add_argument("--level", help="error correction level", default=0.1, type=float)
-parser.add_argument("--video", help="output file for encoded video")
-parser.add_argument("--scale", help="scale of new frames", default=2, type=int)
+parser.add_argument("-i", "--input", help="input file")
+parser.add_argument("-o", "--output", help="output video file")
+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("-m", "--mix", help="mix frames with original video", action="store_true")
 args = parser.parse_args()
 
 
-if args.video:
-    cap = cv2.VideoCapture(args.file)
-    args.height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) / args.scale)
-    args.width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) / args.scale)
-
-# Make corners
 cheight = cwidth = max(args.height // 10, args.width // 10)
-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)))
 midwidth = args.width - 2 * cwidth
 frame_size = args.height * args.width - 4 * cheight * cwidth
 frame_xor = np.arange(frame_size, dtype=np.uint8)
@@ -38,12 +28,18 @@ frame_xor = np.arange(frame_size, dtype=np.uint8)
 # raptorq can add up to 4 extra bytes
 rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4
 
-with open(args.file, "rb") as f:
+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 * 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)))
+
 print("Data length:", len(data))
 print("Packets:", len(packets))
 
@@ -55,38 +51,26 @@ def get_frame():
     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))
     idx = (idx + 1) % len(packets)
     frame = np.concatenate(
         (
             np.concatenate(
-                (
-                    wcorner,
-                    frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
-                    rcorner,
-                ),
+                (wcorner, frame_data[: cheight * midwidth].reshape((cheight, midwidth)), rcorner),
                 axis=1,
             ),
             frame_data[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_data[frame_size - cheight * midwidth :].reshape((cheight, midwidth)), bcorner),
                 axis=1,
             ),
         )
     )
     return np.stack(
-        (
-            (frame & 0b00000111) * 255 // 7,
-            (frame >> 3 & 0b00000111) * 255 // 7,
-            (frame >> 6 & 0b00000011) * 255 // 3,
-        ),
+        ((frame & 0b00000111) * 255 // 7, (frame >> 3 & 0b00000111) * 255 // 7, (frame >> 6) * 255 // 3),
         axis=-1,
     ).astype(np.uint8)
 
@@ -111,30 +95,33 @@ class EncoderWidget(QWidget):
         self.label.setPixmap(pixmap)
 
 
-if args.video:
-    out = cv2.VideoWriter(
-        args.video,
-        cv2.VideoWriter_fourcc(*"mp4v"),
-        cap.get(cv2.CAP_PROP_FPS),
-        (args.scale * args.width, args.scale * args.height),
-    )
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame[: args.scale * cheight, : args.scale * cwidth] = 255
-        frame[: args.scale * cheight, args.scale * (args.width - cwidth) :] = 255
-        frame[args.scale * (args.height - cheight) :, : args.scale * cwidth] = 255
-        frame[
-            args.scale * (args.height - cheight) :, args.scale * (args.width - cwidth) :
-        ] = 255
-        out.write(
-            (
-                frame.astype(np.int64)
-                * np.repeat(np.repeat(get_frame(), args.scale, 0), args.scale, 1)
-                / 255
-            ).astype(np.uint8)
-        )
+if args.output:
+    if args.mix:
+        cap = cv2.VideoCapture(args.input)
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        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))
+        while cap.isOpened():
+            ret, frame = 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(get_frame(), hscale, 0), wscale, 1)).astype(np.uint8),
+                    cv2.COLOR_RGB2BGR,
+                )
+            )
+    else:
+        out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"mp4v"), args.fps, (4* args.width, 4*args.height))
+        for _ in packets:
+            out.write(cv2.cvtColor(np.repeat(np.repeat(get_frame(), 4, 0), 4, 1), cv2.COLOR_RGB2BGR))
 else:
     input("Seizure warning!")
     app = QApplication([])