124 lines
5.4 KiB
Python
124 lines
5.4 KiB
Python
import argparse
|
|
import cv2
|
|
import numpy as np
|
|
from creedsolo import RSCodec
|
|
from raptorq import Encoder
|
|
|
|
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument("-i", "--input", help="input file")
|
|
parser.add_argument("-o", "--output", help="output video file", default="vid.mkv")
|
|
parser.add_argument("-x", "--height", help="grid height", default=100, type=int)
|
|
parser.add_argument("-y", "--width", help="grid width", default=100, type=int)
|
|
parser.add_argument("-l", "--level", help="error correction level", default=0.1, type=float)
|
|
parser.add_argument("-f", "--fps", help="frame rate", default=30, type=int)
|
|
parser.add_argument("-m", "--mix", help="mix frames with original video", action="store_true")
|
|
parser.add_argument("-v", "--version",
|
|
help="0: 10% corners w/ two-sided one-cell padding; 1: 15% corners w/ four-sided 25% padding.",
|
|
default=0, choices=[0, 1], type=int)
|
|
args = parser.parse_args()
|
|
|
|
if args.version == 0:
|
|
cheight = cwidth = max(args.height // 10, args.width // 10)
|
|
elif args.version == 1:
|
|
# # cell borders are 0.0375% of width/height
|
|
# assert args.height * 3 % 80 == args.width * 3 % 80 == 0 # TODO: less strict better ratio
|
|
# cheight = int(args.height * 0.15)
|
|
# cwidth = int(args.width * 0.15)
|
|
|
|
cheight = cwidth = int(max(args.height, args.width) * 0.16)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
midwidth = args.width - 2 * cwidth
|
|
frame_size = args.height * args.width - 4 * cheight * cwidth
|
|
# Divide by 8 / 3 for 3-bit color
|
|
frame_bytes = frame_size * 3 // 8
|
|
frame_xor = np.arange(frame_bytes, dtype=np.uint8)
|
|
# reedsolo breaks message into 255-byte chunks
|
|
# raptorq can add up to 4 extra bytes
|
|
rs_bytes = frame_bytes - (frame_bytes + 254) // 255 * int(args.level * 255) - 4
|
|
|
|
with open(args.input, "rb") as f:
|
|
data = f.read()
|
|
rsc = RSCodec(int(args.level * 255))
|
|
encoder = Encoder.with_defaults(data, rs_bytes)
|
|
packets = encoder.get_encoded_packets(int(len(data) / rs_bytes * (1 / (1 - args.level) - 1)))
|
|
|
|
# Make corners
|
|
if args.version == 0:
|
|
ones = np.ones((cheight - 1, cwidth - 1))
|
|
zeros = np.zeros((cheight - 1, cwidth - 1))
|
|
wcorner = np.pad(np.dstack((ones, ones, ones)), ((0, 1), (0, 1), (0, 0)))
|
|
rcorner = np.pad(np.dstack((ones, zeros, zeros)), ((0, 1), (1, 0), (0, 0)))
|
|
gcorner = np.pad(np.dstack((zeros, ones, zeros)), ((1, 0), (0, 1), (0, 0)))
|
|
bcorner = np.pad(np.dstack((zeros, zeros, ones)), ((1, 0), (1, 0), (0, 0)))
|
|
elif args.version == 1:
|
|
zeros = np.zeros((cheight, cwidth, 3))
|
|
wcorner = zeros.copy()
|
|
rcorner = zeros.copy()
|
|
gcorner = zeros.copy()
|
|
bcorner = zeros.copy()
|
|
black_border_h, black_border_w = cheight // 4, cwidth // 4
|
|
for corner_arr, ones_channel_ind in [(wcorner, 0), (wcorner, 1), (wcorner, 2),
|
|
(rcorner, 0), (gcorner, 1), (bcorner, 2)]:
|
|
corner_arr[black_border_h:-black_border_h, black_border_w:-black_border_w, ones_channel_ind] = np.ones((cheight // 2, cwidth // 2))
|
|
|
|
# Output flags for decoder
|
|
print(f"-x {args.height} -y {args.width} -l {args.level} -s {len(data)} -p {len(packets[0])} -v {args.version}", end="")
|
|
|
|
|
|
def mkframe(packet):
|
|
frame = np.array(rsc.encode(bytearray(packet)))
|
|
frame = np.pad(frame, (0, frame_bytes - len(frame))) ^ frame_xor
|
|
reshape_len = frame_bytes // 255 * 255
|
|
# Space out elements in each size 255 chunk
|
|
frame[:reshape_len] = np.ravel(frame[:reshape_len].reshape(reshape_len // 255, 255), "F")
|
|
frame = np.unpackbits(frame)
|
|
# Pad to be multiple of 3 so we can reshape into RGB channels
|
|
frame = np.pad(frame, (0, (3 - len(frame)) % 3))
|
|
frame = np.reshape(frame, (frame_size, 3))
|
|
frame = np.concatenate(
|
|
(
|
|
np.concatenate(
|
|
(wcorner, frame[: cheight * midwidth].reshape((cheight, midwidth, 3)), rcorner),
|
|
axis=1,
|
|
),
|
|
frame[cheight * midwidth: frame_size - cheight * midwidth].reshape(
|
|
(args.height - 2 * cheight, args.width, 3)
|
|
),
|
|
np.concatenate(
|
|
(gcorner, frame[frame_size - cheight * midwidth:].reshape((cheight, midwidth, 3)), bcorner),
|
|
axis=1,
|
|
),
|
|
)
|
|
)
|
|
return frame.astype(np.uint8) * 255
|
|
|
|
|
|
if args.mix:
|
|
# Mix frames with original video
|
|
cap = cv2.VideoCapture(args.input)
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
hscale = height // args.height
|
|
wscale = width // args.width
|
|
out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (width, height))
|
|
i = 0
|
|
while cap.isOpened():
|
|
ret, vidframe = cap.read()
|
|
if not ret:
|
|
break
|
|
vidframe[: hscale * cheight, : wscale * cwidth] = 0
|
|
vidframe[: hscale * cheight, wscale * (args.width - cwidth) :] = 0
|
|
vidframe[hscale * (args.height - cheight) :, : wscale * cwidth] = 0
|
|
vidframe[hscale * (args.height - cheight) :, wscale * (args.width - cwidth) :] = 0
|
|
frame = np.repeat(np.repeat(mkframe(packets[i]), hscale, 0), wscale, 1)
|
|
# Set edges in original video to black
|
|
frame[vidframe % 255 != 0] = 0
|
|
out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
|
i = (i + 1) % len(packets)
|
|
else:
|
|
# Create a new video
|
|
out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (args.width, args.height))
|
|
for packet in packets:
|
|
out.write(cv2.cvtColor(mkframe(packet), cv2.COLOR_RGB2BGR))
|