Spaces:
Sleeping
Sleeping
File size: 1,243 Bytes
e389f7b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
'''
LastEditors: fcy
'''
import torchac
import torch
import numpy as np
# Encode to bytestream.
seed=6
torch.manual_seed(seed)
np.random.seed(seed)
dim = 500
symsNum = 40000
pdf = np.random.rand(symsNum,dim)
pdf = pdf / (np.sum(pdf,1,keepdims=True))
sym = torch.ShortTensor(np.random.randint(0,dim,symsNum,dtype=np.int16))
def pdf_convert_to_cdf_and_normalize(pdf):
assert pdf.ndim==2
pdf = pdf / (np.sum(pdf,1,keepdims=True))/(1+10**(-10))
cdfF = np.cumsum( pdf, axis=1)
cdfF = np.hstack((np.zeros((pdf.shape[0],1)),cdfF))
return cdfF
output_cdf = torch.Tensor(pdf_convert_to_cdf_and_normalize(pdf)) # Get CDF from your model, shape B, C, H, W, Lp
byte_stream = torchac.encode_float_cdf(output_cdf, sym, check_input_bounds=True)
# pdf = np.diff(cdfF)
# print( -np.log2(pdf[range(0,oct_len),sym]).sum())
# Number of bits taken by the stream
real_bits = len(byte_stream) * 8
print(real_bits)
# Write to a file.
with open('outfile.b', 'wb') as fout:
fout.write(byte_stream)
# Read from a file.
with open('outfile.b', 'rb') as fin:
byte_stream = fin.read()
# Decode from bytestream.
sym_out = torchac.decode_float_cdf(output_cdf, byte_stream)
# Output will be equal to the input.
assert sym_out.equal(sym)
|