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

Downsample+flood fill algorithm (it almost works!)

Browse source
This commit is contained in:
Anthony Wang 2024-04-22 00:13:46 -04:00
parent b5ffcec41a
commit ef73c1183f
Signed by: a
SSH key fingerprint: SHA256:B5ADfMCqd2M7d/jtXDoihAV/yfXOAbWWri9+GdCN4hQ
3 changed files with 156 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

144
decoder.py
View file

@ -1,5 +1,7 @@
import argparse
import collections
import cv2
import matplotlib.pyplot as plt
import numpy as np
from creedsolo import RSCodec
from raptorq import Decoder
@ -14,34 +16,148 @@ parser.add_argument("--level", help="error correction level", default=0.1, type=
parser.add_argument("--device", help="camera device index", default=1, type=int)
args = parser.parse_args()
cheight = args.height // 10
cwidth = args.width // 10
cheight = cwidth = max(args.height // 10, args.width // 10)
frame_size = args.height * args.width - 4 * cheight * cwidth
frame_xor = np.arange(frame_size, dtype=np.uint8)
rs_size = int(frame_size * (1 - args.level))
rs_size = frame_size - int((frame_size + 254) / 255) * int(args.level * 255) - 4
rsc = RSCodec(frame_size - rs_size)
raptor_decoder = Decoder.with_defaults(args.len, rs_size)
rsc = RSCodec(int(args.level * 255))
decoder = Decoder.with_defaults(args.len, rs_size)
data = None
cap = cv2.VideoCapture(args.device)
# cap = cv2.VideoCapture(args.device)
while data is None:
ret, raw_frame = cap.read()
if not ret:
continue
color_frame = decode(raw_frame) # TODO
# ret, raw_frame = cap.read()
# if not ret:
# continue
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()
def max_in_orig(x):
return tuple(
np.array(np.unravel_index(np.argmax(x), x.shape)) * scale + scale // 2
)
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)
# 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
widx, wcol = flood_fill(widx)
ridx, rcol = flood_fill(ridx)
gidx, gcol = flood_fill(gidx)
bidx, bcol = flood_fill(bidx)
# 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)
# 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)
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 & 0b00000111)
+ (color_frame[:, :, 1] >> 2 & 0b00111000)
(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.heigth - cheight, cwidth : args.width - cwidth].flatten(),
frame[args.height - cheight :, cwidth : args.width - cwidth].flatten(),
)
)
data = raptor_decoder.decode(rsc.decode(frame_data ^ frame_xor))
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:
f.write(data)
cap.release()

25
encoder.py
View file

@ -17,8 +17,12 @@ parser.add_argument("--level", help="error correction level", default=0.1, type=
args = parser.parse_args()
cheight = args.height // 10
cwidth = args.width // 10
# 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)
@ -31,6 +35,7 @@ with open(args.file, "rb") as f:
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))
print("Data length:", len(data))
print("Packets:", len(packets))
input("Seizure warning!")
@ -59,9 +64,9 @@ class EncoderWidget(QWidget):
(
np.concatenate(
(
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
wcorner,
frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
rcorner,
),
axis=1,
),
@ -70,21 +75,25 @@ class EncoderWidget(QWidget):
].reshape((args.height - 2 * cheight, args.width)),
np.concatenate(
(
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
gcorner,
frame_data[frame_size - cheight * midwidth :].reshape(
(cheight, midwidth)
),
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
bcorner,
),
axis=1,
),
)
)
color_frame = np.stack(
(frame << 5 & 0b11100000, frame << 2 & 0b11100000, frame & 0b11000000),
(
(frame & 0b00000111) * 255 // 7,
(frame >> 3 & 0b00000111) * 255 // 7,
(frame >> 6 & 0b00000011) * 255 // 3,
),
axis=-1,
)
img = Image.fromarray(color_frame)
img = Image.fromarray(color_frame.astype(np.uint8))
qt_img = ImageQt.ImageQt(img)
pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
self.label.setPixmap(pixmap)

9
requirements.txt
View file

@ -1,9 +1,18 @@
contourpy==1.2.1
cycler==0.12.1
Cython==3.0.10
fonttools==4.51.0
kiwisolver==1.4.5
matplotlib==3.8.4
numpy==1.26.4
opencv-python==4.9.0.80
packaging==24.0
pillow==10.3.0
pyparsing==3.1.2
PyQt6==6.6.1
PyQt6-Qt6==6.6.3
PyQt6-sip==13.6.0
python-dateutil==2.9.0.post0
raptorq==2.0.0
reedsolo==1.7.0
six==1.16.0