k/6.8301-Project
1
1
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

Determining colors by averaging over expected region

Browse source
This commit is contained in:
Kevin Zhao 2024-05-07 22:07:15 -04:00
parent 3a40a8fb03
commit 74dcc919d6
3 changed files with 67 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
README.md
View file

@ -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`

22
decoder.py
View file

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

50
decoding_utils.py
View file

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