Working with 0.1 playback spd, but only 20kbps; commented out debug code

2 files changed, 86 insertions, 20 deletions

diff --git a/decoder.py b/decoder.py
index 1e08c65..e8e548f 100644
--- a/decoder.py
+++ b/decoder.py
@@ -94,6 +94,23 @@ if args.version == 0:
 elif args.version == 1:
     localize_corners = localize_corners_wrapper(args, input_crop_size)
 
+# ####
+# vid_frames = []
+# # cap = cv2.VideoCapture("/Users/kevinzhao/Downloads/IMG_0994.MOV")
+# cap = cv2.VideoCapture("vid_tiny_v1.mkv")
+# data = None
+# start_time = time.time()
+# while data is None:
+#     ret, raw_frame = cap.read()
+#     if not ret:
+#         print("End of stream")
+#         break
+#     vid_frames.append(raw_frame)
+# gtruth = cv2.cvtColor(vid_frames[0], cv2.COLOR_BGR2RGB)
+# ####
+
+
+
 if args.input.isdecimal():
     args.input = int(args.input)
 cap = cv2.VideoCapture(args.input)
@@ -102,6 +119,7 @@ start_time = time.time()
 while data is None:
     try:
         ret, raw_frame = cap.read()
+        # raw_frame = cv2.resize(raw_frame, (1024, 1024), interpolation=cv2.INTER_NEAREST)  # TODO: remove
         if not ret:
             print("End of stream")
             break
@@ -114,7 +132,7 @@ while data is None:
                 raw_frame = raw_frame[X // 4: 3 * X // 4, Y // 4: 3 * Y // 4]
         elif args.version == 1:
             h, w, _ = raw_frame.shape
-            raw_frame = raw_frame[(h - input_crop_size) // 2:-(h - input_crop_size) // 2,
+            raw_frame = raw_frame[(h - input_crop_size) // 2:-(h - input_crop_size) // 2,  # TODO: put back
                             (w - input_crop_size) // 2:-(w - input_crop_size) // 2]
 
         cv2.imshow("", raw_frame)
@@ -133,11 +151,23 @@ while data is None:
 
         # Convert to new color space
         # calibrated_frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 192).astype(np.uint8)
-        calibrated_frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 128).astype(np.uint8)
-        fig, axs = plt.subplots(1, 2)
-        axs[0].imshow(frame)
-        axs[1].imshow(calibrated_frame * 255)
-        plt.show()
+        # calibrated_frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 192).astype(np.uint8) * 255
+        calibrated_frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 150).astype(np.uint8) * 255
+        # fig, axs = plt.subplots(1, 2)
+        # axs[0].imshow(frame)
+        # axs[1].imshow(calibrated_frame)
+        # plt.show()
+
+        # fig, axs = plt.subplots(1, 2)
+        # correct_mask = np.logical_not((calibrated_frame != gtruth).any(axis=2))
+        # calibrated_frame_copy = calibrated_frame.copy()
+        # gtruth_copy = gtruth.copy()
+        # calibrated_frame_copy[correct_mask] = [0, 0, 0]
+        # gtruth_copy[correct_mask] = [0, 0, 0]
+        # axs[0].imshow(gtruth_copy)
+        # axs[1].imshow(calibrated_frame_copy)
+        # plt.show()
+
         calibrated_frame = np.packbits(
             np.concatenate(
                 (
@@ -155,6 +185,7 @@ while data is None:
         break
     except Exception as e:
         print(e)
+
 cap.release()
 with open(args.output, "wb") as f:
     f.write(data)
diff --git a/decoding_utils.py b/decoding_utils.py
index 993e2f9..d3bf2d7 100644
--- a/decoding_utils.py
+++ b/decoding_utils.py
@@ -208,26 +208,61 @@ def localize_corners_wrapper(args, input_crop_size, debug=False):
 
         cheight = int(args.height * 0.15)
         cwidth = int(args.width * 0.15)
-        cch = int(args.height * 0.15) // 4 - 1  # 0-indexed
-        ccw = int(args.width * 0.15) // 4 - 1
+        cch = int(args.height * 0.15) // 4
+        ccw = int(args.width * 0.15) // 4
+
+        # plt.imshow(cropped_frame)
+        # plt.show()
+
+        grid_coords = np.float32([
+            [ccw, cch],
+            [args.width - ccw, cch],
+            [ccw, args.height - cch],
+            [args.width - ccw, args.height - cch],
+        ])
+
+        grid_coords -= 1/2
+
+        #
+        # grid_coords *= orig_w / args.width
+        # torch_frame = transforms_f.perspective(
+        #     transforms.Compose([
+        #         # transforms.Lambda(
+        #         #     lambda img: cv2.resize(img, (stage1_size, stage1_size), interpolation=cv2.INTER_NEAREST)),
+        #         transforms.ToImage(),
+        #         transforms.ToDtype(torch.float32, scale=True),
+        #     ])(cropped_frame),
+        #     orig_pred_pts,
+        #     grid_coords,
+        # )
+        #
+        # torch_frame = cv2.resize(torch_frame.permute(1, 2, 0).numpy(), (args.width, args.height), interpolation=cv2.INTER_AREA)
+        # # torch_frame = cv2.resize(torch_frame.permute(1, 2, 0).numpy(), (args.width, args.height), interpolation=cv2.INTER_NEAREST)
+        #
+        # # torch_frame = transforms_f.resize(torch_frame, [args.height, args.width]).permute(1, 2, 0).numpy()
+        # # torch_frame = torch_frame.permute(1, 2, 0).numpy()
+        # cropped_frame = (torch_frame * 255).astype(np.uint8)
+        # plt.imshow(cropped_frame)
+        #
+        # plt.axis("off")
+        # plt.show()
+
+
+
+        # grid_coords = np.float32([
+        #     [ccw, cch],
+        #     [args.width - ccw - 1, cch],
+        #     [ccw, args.height - cch - 1],
+        #     [args.width - ccw - 1, args.height - cch - 1],
+        # ])
 
         M = cv2.getPerspectiveTransform(
             np.float32(orig_pred_pts),
-            np.float32(
-                [
-                    [ccw, cch],
-                    [args.width - ccw - 1, cch],
-                    [ccw, args.height - cch - 1],
-                    [args.width - ccw - 1, args.height - cch - 1],
-                ]
-            ),
+            grid_coords,
         )
 
         cropped_frame = cv2.warpPerspective(cropped_frame, M, (args.width, args.height))
-
-        # 1-index
-        cch += 1
-        ccw += 1
+        # cropped_frame = cv2.warpPerspective(cropped_frame, M, (args.width, args.height), flags=cv2.INTER_NEAREST)
 
         padding = math.ceil(max(args.height, args.width) / 80)  # arbitrary