Determining colors by averaging over expected region

3 files changed, 67 insertions, 16 deletions

diff --git a/README.md b/README.md
index cafe92a..91dc308 100644
--- a/README.md
+++ b/README.md
@@ -10,10 +10,19 @@ pip install --upgrade reedsolo --no-binary "reedsolo" --no-cache --config-settin
 
 ## Usage
 
-Encode: `python encoder.py -i in`
+### Encoding v0
+Encode: `python encoder.py -v 0 -i in`
 
 Play video (SEIZURE WARNING): `mpv --scale=nearest --fullscreen --loop --no-keepaspect vid.mkv`
 
+### Encoding v1
+Encode: `python encoder.py -v 1 -i in -x 80 -y 80`
+
+Play video (SEIZURE WARNING): `mpv --scale=nearest --fullscreen --loop vid.mkv`
+
+
+### Decoding
+
 Copy the flags printed by the encoder and pass them to the decoder: `python decoder.py FLAGS`
 
 Formatting: `black -l 120 *.py`
diff --git a/decoder.py b/decoder.py
index 280cd36..1e08c65 100644
--- a/decoder.py
+++ b/decoder.py
@@ -1,6 +1,7 @@
 import argparse
 import time
 import cv2
+import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from creedsolo import RSCodec
@@ -129,23 +130,26 @@ while data is None:
         bcol -= origin
         F = 255 * np.linalg.inv(np.stack((rcol, gcol, bcol)).T)
 
+
         # Convert to new color space
-        frame = (np.squeeze(F @ (frame - origin)[..., np.newaxis]) >= 192).astype(np.uint8)
-        import matplotlib.pyplot as plt
-        plt.imshow(frame * 255)
+        # 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()
-        frame = np.packbits(
+        calibrated_frame = np.packbits(
             np.concatenate(
                 (
-                    frame[:cheight, cwidth: args.width - cwidth].flatten(),
-                    frame[cheight: args.height - cheight].flatten(),
-                    frame[args.height - cheight:, cwidth: args.width - cwidth].flatten(),
+                    calibrated_frame[:cheight, cwidth: args.width - cwidth].flatten(),
+                    calibrated_frame[cheight: args.height - cheight].flatten(),
+                    calibrated_frame[args.height - cheight:, cwidth: args.width - cwidth].flatten(),
                 )
             )
         )
         reshape_len = frame_bytes // 255 * 255
-        frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(255, reshape_len // 255), "F")
-        data = decoder.decode(bytes(rsc.decode(bytearray(frame ^ frame_xor))[0][: args.psize]))
+        calibrated_frame[:reshape_len] = np.ravel(calibrated_frame[:reshape_len].reshape(255, reshape_len // 255), "F")
+        data = decoder.decode(bytes(rsc.decode(bytearray(calibrated_frame ^ frame_xor))[0][: args.psize]))
         print("Decoded frame")
     except KeyboardInterrupt:
         break
diff --git a/decoding_utils.py b/decoding_utils.py
index c8f292a..993e2f9 100644
--- a/decoding_utils.py
+++ b/decoding_utils.py
@@ -1,4 +1,5 @@
 import itertools
+import math
 import time
 
 import cv2
@@ -8,6 +9,7 @@ import torch
 import torchvision
 import torchvision.transforms.v2 as transforms
 import torchvision.transforms.v2.functional as transforms_f
+from matplotlib import pyplot as plt
 
 from corner_training.models import QuantizedV2, QuantizedV5
 from corner_training.utils import get_gaussian_filter, get_bounded_slices
@@ -85,6 +87,9 @@ def localize_corners_wrapper(args, input_crop_size, debug=False):
         with torch.no_grad():
             stage1_pred = stage1_model(stage1_img.unsqueeze(0)).squeeze(0)
 
+        # plt.imshow(stage1_pred.detach().cpu())
+        # plt.show()
+
         if debug:
             print(57, time.time() - start_time)
 
@@ -133,13 +138,13 @@ def localize_corners_wrapper(args, input_crop_size, debug=False):
             print(corners_by_quad)
 
         outer_corners = []
-        corner_colors = []  # by center, currently rounding to the pixel in the original image
+        # corner_colors = []  # by center, currently rounding to the pixel in the original image
         origin = (quad_size, quad_size)
         for quad in range(4):  # TODO: consistent (x, y) or (i, j)
             outer_corners.append(max((l2_dist(corner, origin), corner) for corner in corners_by_quad[quad])[1])
-            corner_colors.append(cropped_frame[int((sum(corner[0] for corner in corners_by_quad[quad]) / 4 * upscale_factor)),
-                                     int((sum(corner[1] for corner in corners_by_quad[quad]) / 4 * upscale_factor))]
-                                 .astype(np.float64))
+            # corner_colors.append(cropped_frame[int((sum(corner[0] for corner in corners_by_quad[quad]) / 4 * upscale_factor)),
+            #                          int((sum(corner[1] for corner in corners_by_quad[quad]) / 4 * upscale_factor))]
+            #                      .astype(np.float64))
 
         stage2_imgs = []
 
@@ -197,8 +202,14 @@ def localize_corners_wrapper(args, input_crop_size, debug=False):
         if debug:
             print(142, time.time() - start_time)
 
-        cch = int(args.height * 0.15) / 4 - 1
-        ccw = int(args.width * 0.15) / 4 - 1
+        # plt.imshow(cropped_frame)
+        # plt.scatter(np.array(orig_pred_pts).T[0], np.array(orig_pred_pts).T[1])
+        # plt.show()
+
+        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
 
         M = cv2.getPerspectiveTransform(
             np.float32(orig_pred_pts),
@@ -214,6 +225,33 @@ def localize_corners_wrapper(args, input_crop_size, debug=False):
 
         cropped_frame = cv2.warpPerspective(cropped_frame, M, (args.width, args.height))
 
+        # 1-index
+        cch += 1
+        ccw += 1
+
+        padding = math.ceil(max(args.height, args.width) / 80)  # arbitrary
+
+        # guessing wildly on +/- 1s
+        white_sq = cropped_frame[cch + padding: cheight - cch - padding,
+                                 ccw + padding: cwidth - ccw - padding]
+        red_sq = cropped_frame[cch + padding: cheight - cch - padding,
+                               args.width - cwidth + ccw + padding: args.width - ccw - padding]
+        green_sq = cropped_frame[args.height - cheight + cch + padding: args.height - cch - padding,
+                                 ccw + padding: cwidth - ccw - padding]
+        blue_sq = cropped_frame[args.height - cheight + cch + padding: args.height - cch - padding,
+                                args.width - cwidth + ccw + padding: args.width - ccw - padding]
+
+        corner_colors = [white_sq.mean(axis=(0, 1)), red_sq.mean(axis=(0, 1)),
+                         green_sq.mean(axis=(0, 1)), blue_sq.mean(axis=(0, 1))]
+        # fig, axs = plt.subplots(2, 3)
+        #
+        # axs[0, 2].imshow(cropped_frame)
+        # axs[0, 0].imshow(white_sq)
+        # axs[0, 1].imshow(red_sq)
+        # axs[1, 0].imshow(green_sq)
+        # axs[1, 1].imshow(blue_sq)
+        # plt.show()
+
         return cropped_frame, corner_colors
 
     return localize_corners