mirror of https://github.com/AlfredoSequeida/fvid
Add H265 and Zfec
This commit is contained in:
parent
2726bf9489
commit
c99c8858ee
257
fvid/fvid.py
257
fvid/fvid.py
|
@ -11,6 +11,11 @@ import io
|
||||||
import gzip
|
import gzip
|
||||||
import json
|
import json
|
||||||
import base64
|
import base64
|
||||||
|
import decimal
|
||||||
|
import random
|
||||||
|
import magic
|
||||||
|
|
||||||
|
from zfec import easyfec as ef
|
||||||
|
|
||||||
from cryptography.hazmat.backends import default_backend
|
from cryptography.hazmat.backends import default_backend
|
||||||
from cryptography.hazmat.primitives import hashes
|
from cryptography.hazmat.primitives import hashes
|
||||||
|
@ -18,7 +23,7 @@ from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
||||||
from Crypto.Cipher import AES
|
from Crypto.Cipher import AES
|
||||||
|
|
||||||
try:
|
try:
|
||||||
from fvid_cython import cy_get_bits_from_image as cy_gbfi
|
from fvid_cython import cy_gbfi, cy_gbfi_h265
|
||||||
|
|
||||||
use_cython = True
|
use_cython = True
|
||||||
except (ImportError, ModuleNotFoundError):
|
except (ImportError, ModuleNotFoundError):
|
||||||
|
@ -34,6 +39,15 @@ NOTDEBUG = True
|
||||||
TEMPVIDEO = "_temp.mp4"
|
TEMPVIDEO = "_temp.mp4"
|
||||||
FRAMERATE = "1"
|
FRAMERATE = "1"
|
||||||
|
|
||||||
|
# DO NOT CHANGE: (2, 3-8) works sometimes
|
||||||
|
# this is the most effieicnt by far though
|
||||||
|
KVAL = 4
|
||||||
|
MVAL = 5
|
||||||
|
# this can by ANY integer that is a multiple of (KVAL/MVAL)
|
||||||
|
# but it MUST stay the same between encoding/decoding
|
||||||
|
# reccomended 8-64
|
||||||
|
BLOCK = 16
|
||||||
|
|
||||||
|
|
||||||
class WrongPassword(Exception):
|
class WrongPassword(Exception):
|
||||||
pass
|
pass
|
||||||
|
@ -67,21 +81,51 @@ def get_password(password_provided: str) -> bytes:
|
||||||
key = kdf.derive(password)
|
key = kdf.derive(password)
|
||||||
return key
|
return key
|
||||||
|
|
||||||
|
def encode_zfec(bit_array: BitArray) -> BitArray:
|
||||||
def get_bits_from_file(filepath: str, key: bytes) -> BitArray:
|
global KVAL, MVAL, BLOCK
|
||||||
"""
|
"""
|
||||||
Get/read bits from file
|
Apply Reed-Solomon error correction every byte to maximize retrieval
|
||||||
|
possibility as opposed to applying it to the entire file
|
||||||
|
|
||||||
|
bit_array -- BitArray containing raw file data
|
||||||
|
ecc -- The error correction value to be used (default DEFAULT_ECC)
|
||||||
|
"""
|
||||||
|
|
||||||
|
bits = bit_array.bin
|
||||||
|
|
||||||
|
# split bits into blocks of bits
|
||||||
|
byte_list = split_string_by_n(bits, BLOCK)
|
||||||
|
|
||||||
|
ecc_bytes = ""
|
||||||
|
|
||||||
|
|
||||||
|
print("Applying Zfec Error Correction...")
|
||||||
|
|
||||||
|
encoder = ef.Encoder(KVAL, MVAL)
|
||||||
|
for b in tqdm(byte_list):
|
||||||
|
ecc_bytes += ''.join(map(bytes.decode, encoder.encode(b.encode('utf-8'))))
|
||||||
|
|
||||||
|
return BitArray(bytes=ecc_bytes.encode('utf-8'))
|
||||||
|
|
||||||
|
def get_bits_from_file(
|
||||||
|
filepath: str, key: bytes, zfec: bool
|
||||||
|
) -> BitArray:
|
||||||
|
"""
|
||||||
|
Get/read bits fom file, encrypt data, and zip
|
||||||
|
|
||||||
filepath -- the file to read
|
filepath -- the file to read
|
||||||
key -- key used to encrypt file
|
key -- key used to encrypt file
|
||||||
|
zfec -- if reed solomon should be used to encode bits
|
||||||
"""
|
"""
|
||||||
|
|
||||||
print("Reading file...")
|
print("Reading file...")
|
||||||
|
|
||||||
bitarray = BitArray(filename=filepath)
|
bitarray = BitArray(filename=filepath)
|
||||||
|
|
||||||
# adding a delimiter to know when the file ends to avoid corrupted files
|
if zfec:
|
||||||
# when retrieving
|
bitarray = encode_zfec(bitarray)
|
||||||
|
|
||||||
|
# encrypt data
|
||||||
cipher = AES.new(key, AES.MODE_EAX, nonce=SALT)
|
cipher = AES.new(key, AES.MODE_EAX, nonce=SALT)
|
||||||
ciphertext, tag = cipher.encrypt_and_digest(bitarray.tobytes())
|
ciphertext, tag = cipher.encrypt_and_digest(bitarray.tobytes())
|
||||||
|
|
||||||
|
@ -89,6 +133,7 @@ def get_bits_from_file(filepath: str, key: bytes) -> BitArray:
|
||||||
|
|
||||||
# because json can only serialize strings, the byte objects are encoded
|
# because json can only serialize strings, the byte objects are encoded
|
||||||
# using base64
|
# using base64
|
||||||
|
|
||||||
data_bytes = json.dumps(
|
data_bytes = json.dumps(
|
||||||
{
|
{
|
||||||
"tag": base64.b64encode(tag).decode("utf-8"),
|
"tag": base64.b64encode(tag).decode("utf-8"),
|
||||||
|
@ -97,7 +142,7 @@ def get_bits_from_file(filepath: str, key: bytes) -> BitArray:
|
||||||
}
|
}
|
||||||
).encode("utf-8")
|
).encode("utf-8")
|
||||||
|
|
||||||
print("Zipping...")
|
# print("Zipping...")
|
||||||
|
|
||||||
# zip
|
# zip
|
||||||
out = io.BytesIO()
|
out = io.BytesIO()
|
||||||
|
@ -109,49 +154,64 @@ def get_bits_from_file(filepath: str, key: bytes) -> BitArray:
|
||||||
del bitarray
|
del bitarray
|
||||||
|
|
||||||
bitarray = BitArray(zip)
|
bitarray = BitArray(zip)
|
||||||
|
# bitarray = BitArray(data_bytes)
|
||||||
|
|
||||||
return bitarray.bin
|
return bitarray.bin
|
||||||
|
|
||||||
|
|
||||||
def get_bits_from_image(image: Image) -> str:
|
def get_bits_from_image(image: Image, use_h265: bool) -> str:
|
||||||
"""
|
"""
|
||||||
extract bits from image (frame) pixels
|
extract bits from image (frame) pixels
|
||||||
|
|
||||||
image -- png image file used to extract bits from
|
image -- png image file used to extract bits from
|
||||||
"""
|
"""
|
||||||
|
|
||||||
if use_cython:
|
# use two different functions so we can type pixel correctly
|
||||||
bits = cy_gbfi(image)
|
if use_cython and not use_h265:
|
||||||
return bits
|
return cy_gbfi(image)
|
||||||
|
elif use_cython and use_h265:
|
||||||
|
return cy_gbfi_h265(image)
|
||||||
|
|
||||||
width, height = image.size
|
width, height = image.size
|
||||||
|
|
||||||
px = image.load()
|
px = image.load()
|
||||||
bits = ""
|
bits = ""
|
||||||
|
|
||||||
for y in range(height):
|
# use separate code path so we dont check inside every loop
|
||||||
for x in range(width):
|
if not use_h265:
|
||||||
|
for y in range(height):
|
||||||
|
for x in range(width):
|
||||||
|
pixel = px[x, y]
|
||||||
|
pixel_bin_rep = "0"
|
||||||
|
|
||||||
pixel = px[x, y]
|
# if the white difference is smaller, that means the pixel is
|
||||||
|
# closer to white, otherwise, the pixel must be black
|
||||||
|
if (
|
||||||
|
abs(pixel[0] - 255) < abs(pixel[0] - 0)
|
||||||
|
and abs(pixel[1] - 255) < abs(pixel[1] - 0)
|
||||||
|
and abs(pixel[2] - 255) < abs(pixel[2] - 0)
|
||||||
|
):
|
||||||
|
pixel_bin_rep = "1"
|
||||||
|
|
||||||
pixel_bin_rep = "0"
|
# adding bits
|
||||||
|
bits += pixel_bin_rep
|
||||||
|
else:
|
||||||
|
for y in range(height):
|
||||||
|
for x in range(width):
|
||||||
|
pixel = px[x, y]
|
||||||
|
pixel_bin_rep = "0"
|
||||||
|
|
||||||
# if the white difference is smaller, that means the pixel is
|
# pixel is either 0 or 255, black or white
|
||||||
# closer to white, otherwise, the pixel must be black
|
if pixel == 255:
|
||||||
if (
|
pixel_bin_rep = "1"
|
||||||
abs(pixel[0] - 255) < abs(pixel[0] - 0)
|
|
||||||
and abs(pixel[1] - 255) < abs(pixel[1] - 0)
|
|
||||||
and abs(pixel[2] - 255) < abs(pixel[2] - 0)
|
|
||||||
):
|
|
||||||
pixel_bin_rep = "1"
|
|
||||||
|
|
||||||
# adding bits
|
# adding bits
|
||||||
bits += pixel_bin_rep
|
bits += pixel_bin_rep
|
||||||
|
|
||||||
return bits
|
return bits
|
||||||
|
|
||||||
|
|
||||||
def get_bits_from_video(video_filepath: str) -> str:
|
def get_bits_from_video(video_filepath: str, use_h265: bool) -> str:
|
||||||
"""
|
"""
|
||||||
extract the bits from a video by frame (using a sequence of images)
|
extract the bits from a video by frame (using a sequence of images)
|
||||||
|
|
||||||
|
@ -161,14 +221,24 @@ def get_bits_from_video(video_filepath: str) -> str:
|
||||||
print("Reading video...")
|
print("Reading video...")
|
||||||
|
|
||||||
image_sequence = []
|
image_sequence = []
|
||||||
os.system(
|
if use_h265:
|
||||||
"ffmpeg -i "
|
os.system(
|
||||||
+ video_filepath
|
"ffmpeg -i "
|
||||||
+ " -c:v libx264rgb -filter:v fps=fps="
|
+ video_filepath
|
||||||
+ FRAMERATE
|
+ " -c:v libx265 -filter:v fps=fps="
|
||||||
+ " "
|
+ FRAMERATE
|
||||||
+ TEMPVIDEO
|
+ " -x265-params lossless=1 -preset 6 -tune grain "
|
||||||
)
|
+ TEMPVIDEO
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
os.system(
|
||||||
|
"ffmpeg -i "
|
||||||
|
+ video_filepath
|
||||||
|
+ " -c:v libx264rgb -filter:v fps=fps="
|
||||||
|
+ FRAMERATE
|
||||||
|
+ " "
|
||||||
|
+ TEMPVIDEO
|
||||||
|
)
|
||||||
os.system(
|
os.system(
|
||||||
"ffmpeg -i " + TEMPVIDEO + " ./fvid_frames/decoded_frames_%d.png"
|
"ffmpeg -i " + TEMPVIDEO + " ./fvid_frames/decoded_frames_%d.png"
|
||||||
)
|
)
|
||||||
|
@ -185,19 +255,59 @@ def get_bits_from_video(video_filepath: str) -> str:
|
||||||
|
|
||||||
if use_cython:
|
if use_cython:
|
||||||
print("Using Cython...")
|
print("Using Cython...")
|
||||||
|
|
||||||
for index in tqdm(range(sequence_length)):
|
for index in tqdm(range(sequence_length)):
|
||||||
bits += get_bits_from_image(image_sequence[index])
|
bits += get_bits_from_image(image_sequence[index], use_h265)
|
||||||
|
|
||||||
return bits
|
return bits
|
||||||
|
|
||||||
|
|
||||||
def save_bits_to_file(file_path: str, bits: str, key: bytes):
|
def decode_zfec(data_bytes: bytes) -> bytes:
|
||||||
|
global KVAL, MVAL, BLOCK
|
||||||
|
|
||||||
|
byte_list = split_string_by_n(data_bytes, int(BLOCK*(MVAL/KVAL)))
|
||||||
|
|
||||||
|
# appending to a single bytes object is very slow so we make 50-51 and combine at the end
|
||||||
|
decoded_bytes = [bytes()] * (int(len(byte_list) / 50) + 1)
|
||||||
|
|
||||||
|
print("Decoding Zfec Error Correction...")
|
||||||
|
|
||||||
|
decoder = ef.Decoder(KVAL, MVAL)
|
||||||
|
i = 0
|
||||||
|
for b in tqdm(byte_list):
|
||||||
|
base = split_string_by_n(b, len(b) // MVAL)
|
||||||
|
decoded_str_1 = decoder.decode(base[:KVAL], list(range(KVAL)), 0)
|
||||||
|
decoded_str_2 = decoder.decode(base[1:KVAL+1], list(range(KVAL+1))[1:], 0)
|
||||||
|
if decoded_str_1 == decoded_str_2:
|
||||||
|
decoded_bytes[i//50] += decoded_str_1
|
||||||
|
else: # its corrupted here
|
||||||
|
j = 10
|
||||||
|
while j > 0 and decoded_str_1 != decoded_str_2:
|
||||||
|
random.shuffle(base)
|
||||||
|
decoded_str_1 = decoder.decode(base[:KVAL], list(range(KVAL)), 0)
|
||||||
|
decoded_str_2 = decoder.decode(base[1:KVAL+1], list(range(KVAL+1))[1:], 0)
|
||||||
|
j -= 1
|
||||||
|
decoded_bytes[i//50] += decoded_str_1 # it should be correct by now
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
decoded_bytestring = bytes()
|
||||||
|
|
||||||
|
for bytestring in tqdm(decoded_bytes):
|
||||||
|
decoded_bytestring += bytestring
|
||||||
|
|
||||||
|
return decoded_bytestring
|
||||||
|
|
||||||
|
|
||||||
|
def save_bits_to_file(
|
||||||
|
file_path: str, bits: str, key: bytes, zfec: bool
|
||||||
|
):
|
||||||
"""
|
"""
|
||||||
save/write bits to a file
|
save/write bits to a file
|
||||||
|
|
||||||
file_path -- the path to write to
|
file_path -- the path to write to
|
||||||
bits -- the bits to write
|
bits -- the bits to write
|
||||||
key -- key userd for file decryption
|
key -- key userd for file decryption
|
||||||
|
zfec -- needed if reed solomon was used to encode bits
|
||||||
"""
|
"""
|
||||||
|
|
||||||
bitstring = Bits(bin=bits)
|
bitstring = Bits(bin=bits)
|
||||||
|
@ -207,6 +317,9 @@ def save_bits_to_file(file_path: str, bits: str, key: bytes):
|
||||||
in_ = io.BytesIO()
|
in_ = io.BytesIO()
|
||||||
in_.write(bitstring.bytes)
|
in_.write(bitstring.bytes)
|
||||||
in_.seek(0)
|
in_.seek(0)
|
||||||
|
# DOES NOT WORK IF WE DONT CHECK BUFFER, UNSURE WHY
|
||||||
|
filetype = magic.from_buffer(in_.read())
|
||||||
|
in_.seek(0)
|
||||||
with gzip.GzipFile(fileobj=in_, mode="rb") as fo:
|
with gzip.GzipFile(fileobj=in_, mode="rb") as fo:
|
||||||
bitstring = fo.read()
|
bitstring = fo.read()
|
||||||
# zip
|
# zip
|
||||||
|
@ -220,8 +333,9 @@ def save_bits_to_file(file_path: str, bits: str, key: bytes):
|
||||||
|
|
||||||
filename = data["filename"]
|
filename = data["filename"]
|
||||||
|
|
||||||
|
# decrypting data
|
||||||
cipher = AES.new(key, AES.MODE_EAX, nonce=SALT)
|
cipher = AES.new(key, AES.MODE_EAX, nonce=SALT)
|
||||||
bitstring = cipher.decrypt(ciphertext)
|
data_bytes = cipher.decrypt(ciphertext)
|
||||||
|
|
||||||
print("Checking integrity...")
|
print("Checking integrity...")
|
||||||
|
|
||||||
|
@ -230,7 +344,12 @@ def save_bits_to_file(file_path: str, bits: str, key: bytes):
|
||||||
except ValueError:
|
except ValueError:
|
||||||
raise WrongPassword("Key incorrect or message corrupted")
|
raise WrongPassword("Key incorrect or message corrupted")
|
||||||
|
|
||||||
bitstring = BitArray(bitstring)
|
bitstring = Bits(data_bytes)
|
||||||
|
|
||||||
|
if zfec:
|
||||||
|
bitstring = Bits(
|
||||||
|
"0b" + decode_zfec(data_bytes).decode("utf-8")
|
||||||
|
)
|
||||||
|
|
||||||
# If filepath not passed in use default otherwise used passed in filepath
|
# If filepath not passed in use default otherwise used passed in filepath
|
||||||
if file_path == None:
|
if file_path == None:
|
||||||
|
@ -286,11 +405,12 @@ def make_image_sequence(bitstring: BitArray, resolution: tuple = (1920, 1080)):
|
||||||
|
|
||||||
index = 1
|
index = 1
|
||||||
bitlist = bitlist[::-1]
|
bitlist = bitlist[::-1]
|
||||||
|
|
||||||
print("Saving frames...")
|
print("Saving frames...")
|
||||||
|
|
||||||
for _ in tqdm(range(len(bitlist))):
|
for _ in tqdm(range(len(bitlist))):
|
||||||
bitl = bitlist.pop()
|
bitl = bitlist.pop()
|
||||||
image_bits = list(map(int, bitl))
|
image_bits = list(map(int, bitl))
|
||||||
# print(image_bits)
|
|
||||||
|
|
||||||
image = Image.new("1", (width, height))
|
image = Image.new("1", (width, height))
|
||||||
image.putdata(image_bits)
|
image.putdata(image_bits)
|
||||||
|
@ -298,7 +418,7 @@ def make_image_sequence(bitstring: BitArray, resolution: tuple = (1920, 1080)):
|
||||||
index += 1
|
index += 1
|
||||||
|
|
||||||
|
|
||||||
def make_video(output_filepath: str, framerate: int = FRAMERATE):
|
def make_video(output_filepath: str, framerate: int = FRAMERATE, use_h265: bool = False):
|
||||||
"""
|
"""
|
||||||
Create video using ffmpeg
|
Create video using ffmpeg
|
||||||
|
|
||||||
|
@ -311,13 +431,21 @@ def make_video(output_filepath: str, framerate: int = FRAMERATE):
|
||||||
else:
|
else:
|
||||||
outputfile = output_filepath
|
outputfile = output_filepath
|
||||||
|
|
||||||
os.system(
|
if use_h265:
|
||||||
"ffmpeg -r "
|
os.system(
|
||||||
+ framerate
|
"ffmpeg -r "
|
||||||
+ " -i ./fvid_frames/encoded_frames_%d.png -c:v libx264rgb "
|
+ framerate
|
||||||
+ outputfile
|
+ " -i ./fvid_frames/encoded_frames_%d.png -c:v libx265 "
|
||||||
)
|
+ " -x265-params lossless=1 -preset 6 -tune grain "
|
||||||
|
+ outputfile
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
os.system(
|
||||||
|
"ffmpeg -r "
|
||||||
|
+ framerate
|
||||||
|
+ " -i ./fvid_frames/encoded_frames_%d.png -c:v libx264rgb "
|
||||||
|
+ outputfile
|
||||||
|
)
|
||||||
|
|
||||||
def cleanup():
|
def cleanup():
|
||||||
"""
|
"""
|
||||||
|
@ -336,7 +464,6 @@ def setup():
|
||||||
if not os.path.exists(FRAMES_DIR):
|
if not os.path.exists(FRAMES_DIR):
|
||||||
os.makedirs(FRAMES_DIR)
|
os.makedirs(FRAMES_DIR)
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
global FRAMERATE
|
global FRAMERATE
|
||||||
parser = argparse.ArgumentParser(description="save files as videos")
|
parser = argparse.ArgumentParser(description="save files as videos")
|
||||||
|
@ -364,6 +491,27 @@ def main():
|
||||||
type=str,
|
type=str,
|
||||||
default="default",
|
default="default",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"-z",
|
||||||
|
"--zfec",
|
||||||
|
help=(
|
||||||
|
"Apply Zfec error correcting. This is helpful if you're"
|
||||||
|
" finding that your data is not being decoded correctly. It adds"
|
||||||
|
" 2 extra bits per byte making it possible to recover all 8 bits"
|
||||||
|
" in the case the data changes during the decoding process at"
|
||||||
|
" the cost of making your video files larger. Note, if you use"
|
||||||
|
" this option, you must also use the -r flag to decode a video"
|
||||||
|
" back to a file, otherwise, your data will not be recovered"
|
||||||
|
" correctly."
|
||||||
|
),
|
||||||
|
action="store_true",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"-5",
|
||||||
|
"--h265",
|
||||||
|
help="Use H.265 codec for improved efficiency",
|
||||||
|
action="store_true",
|
||||||
|
)
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
@ -391,14 +539,14 @@ def main():
|
||||||
key = get_password(args.password)
|
key = get_password(args.password)
|
||||||
|
|
||||||
if args.decode:
|
if args.decode:
|
||||||
bits = get_bits_from_video(args.input)
|
bits = get_bits_from_video(args.input, args.h265)
|
||||||
|
|
||||||
file_path = None
|
file_path = None
|
||||||
|
|
||||||
if args.output:
|
if args.output:
|
||||||
file_path = args.output
|
file_path = args.output
|
||||||
|
|
||||||
save_bits_to_file(file_path, bits, key)
|
save_bits_to_file(file_path, bits, key, args.zfec)
|
||||||
|
|
||||||
elif args.encode:
|
elif args.encode:
|
||||||
|
|
||||||
|
@ -415,7 +563,7 @@ def main():
|
||||||
)
|
)
|
||||||
|
|
||||||
# get bits from file
|
# get bits from file
|
||||||
bits = get_bits_from_file(args.input, key)
|
bits = get_bits_from_file(args.input, key, args.zfec)
|
||||||
|
|
||||||
# create image sequence
|
# create image sequence
|
||||||
make_image_sequence(bits)
|
make_image_sequence(bits)
|
||||||
|
@ -425,6 +573,9 @@ def main():
|
||||||
if args.output:
|
if args.output:
|
||||||
video_file_path = args.output
|
video_file_path = args.output
|
||||||
|
|
||||||
make_video(video_file_path, args.framerate)
|
make_video(video_file_path, args.framerate, args.h265)
|
||||||
|
|
||||||
cleanup()
|
cleanup()
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
Loading…
Reference in New Issue