Implement ECC, webcam capture

This commit is contained in:
Anthony Wang 2024-04-21 01:15:52 -04:00
parent 762d99dc02
commit cef9bab49a
Signed by: a
SSH key fingerprint: SHA256:B5ADfMCqd2M7d/jtXDoihAV/yfXOAbWWri9+GdCN4hQ
3 changed files with 108 additions and 17 deletions

40
decoder.py Normal file
View file

@ -0,0 +1,40 @@
import argparse
import cv2
import numpy as np
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)
args = parser.parse_args()
cheight = args.height // 10
cwidth = args.width // 10
frame_size = args.height * args.width - 4 * cheight * cwidth
rs_size = int(frame_size * (1 - args.level))
rsc = RSCodec(frame_size - rs_size)
raptor_decoder = Decoder.with_defaults(args.size, rs_size)
data = None
cap = cv2.VideoCapture(0)
while data is None:
ret, frame = cap.read()
if not ret:
continue
frame_full = decode(frame) # TODO
frame_data = np.concatenate(
(
frame_full[:cheight, cwidth : args.width - cwidth].flatten(),
frame_full[cheight : args.height - cheight].flatten(),
frame_full[args.heigth - cheight, cwidth : args.width - cwidth].flatten(),
)
)
data = raptor_decoder.decode(rsc.decode(frame_data))
with open(args.file, "wb") as f:
f.write(data)
cap.release()

View file

@ -1,14 +1,38 @@
import argparse
import sys
import numpy as np
from creedsolo import RSCodec
from PyQt6.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout
from PyQt6.QtGui import QPixmap
from PyQt6.QtCore import QTimer
from PIL import Image, ImageQt
from raptorq import Encoder
fps = 30
h = 200
w = 200
c = 0
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("file", help="output file for decoded data")
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)
args = parser.parse_args()
cheight = args.height // 10
cwidth = args.width // 10
midwidth = args.width - 2 * cwidth
frame_size = args.height * args.width - 4 * cheight * cwidth
# reedsolo breaks message into 255-byte chunks
# 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:
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))
print("Data length:", len(data))
print("Packets:", len(packets))
input("Seizure warning!")
class EncoderWidget(QWidget):
@ -16,7 +40,8 @@ class EncoderWidget(QWidget):
super().__init__()
self.timer = QTimer(self)
self.timer.timeout.connect(self.update)
self.timer.start(1000 // fps)
self.timer.start(1000 // args.fps)
self.idx = 0
self.label = QLabel(self)
layout = QVBoxLayout(self)
layout.addWidget(self.label)
@ -25,20 +50,45 @@ class EncoderWidget(QWidget):
self.showFullScreen()
def update(self):
global c
if c == 0:
array = np.random.randint(0, 256, (h, w, 3))
c = 1
else:
array = np.zeros((h, w, 3))
c = 0
img = Image.fromarray(array.astype(np.uint8), mode="RGB")
frame_data = np.array(rsc.encode(packets[self.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)))
self.idx = (self.idx + 1) % len(packets)
frame = np.concatenate(
(
np.concatenate(
(
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
),
axis=1,
),
frame_data[
cheight * midwidth : frame_size - cheight * midwidth
].reshape((args.height - 2 * cheight, args.width)),
np.concatenate(
(
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
frame_data[frame_size - cheight * midwidth :].reshape(
(cheight, midwidth)
),
np.zeros((cheight, cwidth), dtype=np.uint8), # TODO
),
axis=1,
),
)
)
color_frame = np.stack(
((frame & 0x03) << 6, (frame >> 2 & 0x07) << 5, (frame >> 5 & 0x7) << 5),
axis=-1,
)
img = Image.fromarray(color_frame)
qt_img = ImageQt.ImageQt(img)
pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
self.label.setPixmap(pixmap)
if __name__ == "__main__":
app = QApplication(sys.argv)
widget = EncoderWidget()
sys.exit(app.exec())
app = QApplication([])
widget = EncoderWidget()
sys.exit(app.exec())

View file

@ -1,5 +1,6 @@
Cython==3.0.10
numpy==1.26.4
opencv-python==4.9.0.80
pillow==10.3.0
PyQt6==6.6.1
PyQt6-Qt6==6.6.3