def define_architecture():
if cuda:
encoder = network.EncoderCell(3).cuda()
binarizer = network.Binarizer(32).cuda()
decoder = network.DecoderCell(32, 3).cuda()
return encoder, binarizer, decoder
def get_models(args, v_compress, bits, encoder_fuse_level, decoder_fuse_level):
"""build the pytorch network architecture for VCII
Arguments:
args {argparse} -- Arguments defined for the network
v_compress {boolean} -- [description]
bits {[type]} -- [description]
encoder_fuse_level {[type]} -- [description]
decoder_fuse_level {[type]} -- [description]
Returns:
encoder, binarizer, decoder, unet -- [description]
"""
encoder = network.EncoderCell(
v_compress=v_compress,
stack=args.stack,
fuse_encoder=args.fuse_encoder,
fuse_level=encoder_fuse_level
).cuda()
binarizer = network.Binarizer(bits).cuda()
decoder = network.DecoderCell(
v_compress=v_compress, shrink=args.shrink,
bits=bits,
fuse_level=decoder_fuse_level
).cuda()
if v_compress:
unet = UNet(3, args.shrink).cuda()
else:
unet = None
return encoder, binarizer, decoder, unet
def get_models(args, v_compress, bits, encoder_fuse_level, decoder_fuse_level,num_vids):
encoder = network.EncoderCell(
v_compress=v_compress,
stack=args.stack,
fuse_encoder=args.fuse_encoder,
fuse_level=encoder_fuse_level
).cuda()
binarizer = network.Binarizer(bits).cuda()
decoder = network.DecoderCell(
v_compress=v_compress, shrink=args.shrink,
bits=bits,
fuse_level=decoder_fuse_level
).cuda()
if v_compress:
unet = UNet(3, args.shrink).cuda()
else:
unet = None
hypernet = network.HyperNetwork(num_vids).cuda()
print("hypernet ",count_parameters(hypernet))
print("encoder ",count_parameters(encoder))
print("decoder ",count_parameters(decoder))
print("binarizer ",count_parameters(binarizer))
print("unet ",count_parameters(unet))
return encoder, binarizer, decoder, unet, hypernet
def define_architecture():
if cuda:
encoder = network.EncoderCell().cuda()
binarizer = network.Binarizer().cuda()
decoder = network.DecoderCell().cuda()
gain = network.GainFactor().cuda()
return encoder, binarizer, decoder, gain
def get_models(args, v_compress, bits, encoder_fuse_level, decoder_fuse_level):
encoder = network.EncoderCell(v_compress=v_compress,
stack=args.stack,
fuse_encoder=args.fuse_encoder,
fuse_level=encoder_fuse_level).cuda()
binarizer = network.Binarizer(bits).cuda()
decoder = network.DecoderCell(v_compress=v_compress,
shrink=args.shrink,
bits=bits,
fuse_level=decoder_fuse_level).cuda()
if v_compress:
unet = UNet(3, args.shrink).cuda()
else:
unet = None
return encoder, binarizer, decoder, unet
train_loader = data.DataLoader(dataset=train_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=1)
print('total images: {}; total batches: {}'.format(len(train_set),
len(train_loader)))
# >>> 512*18*22
# 202752
## load networks on GPU
import network
hypernet = network.HyperNetwork(train_set.vid_count).cuda()
encoder = network.EncoderCell().cuda()
binarizer = network.Binarizer().cuda()
decoder = network.DecoderCell().cuda()
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
print("hypernet ", count_parameters(hypernet))
print("encoder ", count_parameters(encoder))
print("decoder ", count_parameters(decoder))
print("binarizer ", count_parameters(binarizer))
#ff = Feedforward(202752).cuda()
#unet = UNet(9,1).cuda()
solver = optim.Adam([{
image = imread(args.input, mode='RGB')
image = torch.from_numpy(
np.expand_dims(np.transpose(image.astype(np.float32) / 255.0, (2, 0, 1)),
0))
batch_size, input_channels, height, width = image.size()
assert height % 32 == 0 and width % 32 == 0
image = Variable(image, volatile=True)
import network
encoder = network.EncoderCell()
binarizer = network.Binarizer()
decoder = network.DecoderCell()
encoder.eval()
binarizer.eval()
decoder.eval()
encoder.load_state_dict(torch.load(args.model))
binarizer.load_state_dict(
torch.load(args.model.replace('encoder', 'binarizer')))
decoder.load_state_dict(torch.load(args.model.replace('encoder', 'decoder')))
'''
if args.cuda:
encoder = encoder.cuda()
binarizer = binarizer.cuda()
decoder = decoder.cuda()
def ed_process(self,my_object):
bpp=[]
code_img = []
th =self.target
train_loader = torch.utils.data.DataLoader(self.list_patches,batch_size = self.batch_size,num_workers=2,
shuffle=False)
num_batch = len(train_loader)
encoder = network.EncoderCell(3).eval()
binarizer = network.Binarizer(canais).eval()
decoder = network.DecoderCell(canais,3).eval()
encoder.load_state_dict(torch.load(self.path_model,map_location=lambda storage, loc: storage ) )
binarizer.load_state_dict(torch.load(self.path_model.replace('encoder', 'binarizer'),map_location=lambda storage, loc: storage ))
decoder.load_state_dict(torch.load(self.path_model.replace('encoder', 'decoder'),map_location=lambda storage, loc: storage))
if cuda:
encoder = encoder.cuda()
binarizer = binarizer.cuda()
decoder = decoder.cuda()
qbits = np.zeros(1)
list_patches_recons = np.zeros((self.batch_size*num_batch,self.height,self.width,3),dtype='uint8')
for batch,data in enumerate(train_loader):
image = data
if cuda:
image = image.cuda()
#assert height % 32 == 0 and width % 32 == 0
with torch.no_grad():
encoder_h_1 = ((torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)),
(torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)))
encoder_h_2 = ((torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)),
(torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)))
encoder_h_3 = ((torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)),
(torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)))
decoder_h_1 = ((torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)),
(torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)))
decoder_h_2 = ((torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)),
(torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)))
decoder_h_3 = ((torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)),
(torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)))
decoder_h_4 = ((torch.zeros(self.batch_size, 128, self.height // 2, self.width // 2)),
(torch.zeros(self.batch_size, 128, self.height // 2, self.width // 2)))
encoder_h_1 = (encoder_h_1[0], encoder_h_1[1])
encoder_h_2 = (encoder_h_2[0], encoder_h_2[1])
encoder_h_3 = (encoder_h_3[0], encoder_h_3[1])
decoder_h_1 = (decoder_h_1[0], decoder_h_1[1])
decoder_h_2 = (decoder_h_2[0], decoder_h_2[1])
decoder_h_3 = (decoder_h_3[0], decoder_h_3[1])
decoder_h_4 = (decoder_h_4[0], decoder_h_4[1])
if cuda:
encoder_h_1 = (encoder_h_1[0].cuda(), encoder_h_1[1].cuda())
encoder_h_2 = (encoder_h_2[0].cuda(), encoder_h_2[1].cuda())
encoder_h_3 = (encoder_h_3[0].cuda(), encoder_h_3[1].cuda())
decoder_h_1 = (decoder_h_1[0].cuda(), decoder_h_1[1].cuda())
decoder_h_2 = (decoder_h_2[0].cuda(), decoder_h_2[1].cuda())
decoder_h_3 = (decoder_h_3[0].cuda(), decoder_h_3[1].cuda())
decoder_h_4 = (decoder_h_4[0].cuda(), decoder_h_4[1].cuda())
im = torch.zeros(batch_size, 3, self.height, self.width)
res = (image - offset).float()
codes = []
for iters in range(self.qiters):
encoded, encoder_h_1, encoder_h_2, encoder_h_3 = encoder(res, encoder_h_1, encoder_h_2, encoder_h_3)
code = binarizer(encoded)
output, decoder_h_1, decoder_h_2, decoder_h_3, decoder_h_4 = decoder(code, decoder_h_1, decoder_h_2, decoder_h_3, decoder_h_4)
im += output.data.cpu()
res = res - output
codes.append(code.data.cpu().numpy())
qbits += (self.height/16*self.width/16 *canais)*self.batch_size
if self.op_bit_allocation and iters+1>=self.min_iters:
result =[]
for i in range(self.batch_size):
patch = im[i]+offset
patch = np.squeeze(patch.numpy().clip(0, 1) * 255.0).astype(np.uint8).transpose(1, 2, 0)
list_patches_recons[i+batch*self.batch_size,:,:,:] = patch
patch_test = patch
patch_test = Image.fromarray(patch_test, mode = self.color_mode)
patch_test = patch_test.convert('RGB')
patch_ref = (patches[batch*batch_size+i]* 255.0).astype(np.uint8)
patch_ref = Image.fromarray(patch_ref, mode = self.color_mode)
patch_ref = patch_ref.convert('RGB')
result.append(calc_metric(patch_ref,patch_test,'psnr'))
lower = min(result)
media = np.mean(result)
if media>=th and lower>=th*0.97:
print('iteração',iters+1,'PSNR médio', media)
break
if not self.op_bit_allocation:
for i in range(self.batch_size):
patch = im[i] + offset
patch = np.squeeze(patch.numpy().clip(0, 1) * 255.0).astype(np.uint8).transpose(1, 2, 0)
list_patches_recons[i+batch*self.batch_size,:,:,:] = patch
codes = np.asarray(codes).flatten()
if batch==0:
code_img = codes
else:
code_img = np.concatenate((code_img, codes), axis=0)
codes=[]
code_img = list(map(int, code_img))
latent = list(map(lambda x: x != -1, code_img))
size_in = bitstring.Bits(latent).length
compressed = gzip.compress(bitstring.Bits(latent).bytes)
size_out = bitstring.Bits(compressed).length
bpp2 = size_out/(512*768)
bpp = qbits/(768*512)
return list_patches_recons,bpp,bpp2
def __init__(self,
original_path,
model_path,
codes_path,
save_path,
decode_only=False,
gpu=True):
self.gpu = gpu
global encoder
encoder = network.EncoderCell()
global binarizer
binarizer = network.Binarizer()
global decoder
decoder = network.DecoderCell()
if gpu:
self.encoder = encoder.cuda()
bianrizer = binarizer.cuda()
decoder = decoder.cuda()
encoder = encoder.eval()
bianrizer = binarizer.eval()
decoder = decoder.eval()
self.original = original_path
self.codes_path = codes_path
os.system('mkdir -p {}'.format(self.codes_path))
self.save_path = save_path
os.system('mkdir -p {}'.format(self.save_path))
model = model_path + '/encoder.pth'
if gpu:
if decode_only == False:
encoder.load_state_dict(torch.load(model))
binarizer.load_state_dict(
torch.load(model.replace('encoder', 'binarizer')))
decoder.load_state_dict(
torch.load(model.replace('encoder', 'decoder')))
else:
if decode_only == False:
encoder.load_state_dict(
torch.load(model,
map_location=lambda storage, loc: storage))
binarizer.load_state_dict(
torch.load(model.replace('encoder', 'binarizer'),
map_location=lambda storage, loc: storage))
decoder.load_state_dict(
torch.load(model.replace('encoder', 'decoder'),
map_location=lambda storage, loc: storage))
def ed_process(self,my_object):
bpp=[]
#for th in self.threshold:
#for q in range(self.qiters, self.qiters+1):
th =self.target
train_loader = torch.utils.data.DataLoader(self.list_patches,batch_size = self.batch_size,num_workers=2,
shuffle=False)
num_batch = len(train_loader)
encoder = network.EncoderCell(1)
binarizer = network.Binarizer(28)
decoder = network.DecoderCell(28,1)
encoder2 = network.EncoderCell(2)
binarizer2 = network.Binarizer(4)
decoder2 = network.DecoderCell(4,2)
encoder.load_state_dict(torch.load(self.path_model,map_location=lambda storage, loc: storage ) )
binarizer.load_state_dict(torch.load(self.path_model.replace('encoder', 'binarizer'),map_location=lambda storage, loc: storage ))
decoder.load_state_dict(torch.load(self.path_model.replace('encoder', 'decoder'),map_location=lambda storage, loc: storage))
encoder2.load_state_dict(torch.load(self.path_model2,map_location=lambda storage, loc: storage ) )
binarizer2.load_state_dict(torch.load(self.path_model2.replace('encoder', 'binarizer'),map_location=lambda storage, loc: storage ))
decoder2.load_state_dict(torch.load(self.path_model2.replace('encoder', 'decoder'),map_location=lambda storage, loc: storage))
if cuda:
encoder = encoder.cuda()
binarizer = binarizer.cuda()
decoder = decoder.cuda()
encoder2 = encoder2.cuda()
binarizer2 = binarizer2.cuda()
decoder2 = decoder2.cuda()
qbits=np.zeros(1)
list_patches_recons = np.zeros((self.batch_size*num_batch,self.height,self.width,3),dtype='uint8')
print('Batch size',self.batch_size,'Número de batches: ',num_batch,', Dimensão patch',self.h,'x',self.w)
code_img=[]
code_img2 =[]
for batch,data in enumerate(train_loader):
image = data
if cuda:
image = image.cuda()
codes_final=[]
#assert height % 32 == 0 and width % 32 == 0
with torch.no_grad():
image = Variable(image)
encoder_h_1 = ((torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)),
(torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)))
encoder_h_2 = ((torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)),
(torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)))
encoder_h_3 = ((torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)),
(torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)))
decoder_h_1 = ((torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)),
(torch.zeros(self.batch_size, 512, self.height // 16, self.width // 16)))
decoder_h_2 = ((torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)),
(torch.zeros(self.batch_size, 512, self.height // 8, self.width // 8)))
decoder_h_3 = ((torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)),
(torch.zeros(self.batch_size, 256, self.height // 4, self.width // 4)))
decoder_h_4 = ((torch.zeros(self.batch_size, 128, self.height // 2, self.width // 2)),
(torch.zeros(self.batch_size, 128, self.height // 2, self.width // 2)))
encoder_h_1 = (encoder_h_1[0], encoder_h_1[1])
encoder_h_2 = (encoder_h_2[0], encoder_h_2[1])
encoder_h_3 = (encoder_h_3[0], encoder_h_3[1])
decoder_h_1 = (decoder_h_1[0], decoder_h_1[1])
decoder_h_2 = (decoder_h_2[0], decoder_h_2[1])
decoder_h_3 = (decoder_h_3[0], decoder_h_3[1])
decoder_h_4 = (decoder_h_4[0], decoder_h_4[1])
encoder_h_1_2 = (encoder_h_1[0], encoder_h_1[1])
encoder_h_2_2 = (encoder_h_2[0], encoder_h_2[1])
encoder_h_3_2 = (encoder_h_3[0], encoder_h_3[1])
decoder_h_1_2 = (decoder_h_1[0], decoder_h_1[1])
decoder_h_2_2 = (decoder_h_2[0], decoder_h_2[1])
decoder_h_3_2 = (decoder_h_3[0], decoder_h_3[1])
decoder_h_4_2 = (decoder_h_4[0], decoder_h_4[1])
if cuda:
encoder_h_1 = (encoder_h_1[0].cuda(), encoder_h_1[1].cuda())
encoder_h_2 = (encoder_h_2[0].cuda(), encoder_h_2[1].cuda())
encoder_h_3 = (encoder_h_3[0].cuda(), encoder_h_3[1].cuda())
decoder_h_1 = (decoder_h_1[0].cuda(), decoder_h_1[1].cuda())
decoder_h_2 = (decoder_h_2[0].cuda(), decoder_h_2[1].cuda())
decoder_h_3 = (decoder_h_3[0].cuda(), decoder_h_3[1].cuda())
decoder_h_4 = (decoder_h_4[0].cuda(), decoder_h_4[1].cuda())
encoder_h_1_2 = (encoder_h_1[0].cuda(), encoder_h_1[1].cuda())
encoder_h_2_2 = (encoder_h_2[0].cuda(), encoder_h_2[1].cuda())
encoder_h_3_2 = (encoder_h_3[0].cuda(), encoder_h_3[1].cuda())
decoder_h_1_2 = (decoder_h_1[0].cuda(), decoder_h_1[1].cuda())
decoder_h_2_2 = (decoder_h_2[0].cuda(), decoder_h_2[1].cuda())
decoder_h_3_2 = (decoder_h_3[0].cuda(), decoder_h_3[1].cuda())
decoder_h_4_2 = (decoder_h_4[0].cuda(), decoder_h_4[1].cuda())
image = image - 0.5
image = image.float()
im = torch.zeros(batch_size, 1, self.height, self.width) + 0.5
im2 = torch.zeros(batch_size, 2, self.height, self.width) + 0.5
res = torch.zeros(batch_size, 1, self.height, self.width,dtype=torch.float).cuda()
res[:,0,:,:] = image[:,0,:,:]
res2 = image[:,1:,:,:]
result = 0
codes = []
codes2 = []
for iters in range(self.qiters):
encoded, encoder_h_1, encoder_h_2, encoder_h_3 = encoder(res, encoder_h_1, encoder_h_2, encoder_h_3)
code = binarizer(encoded)
output, decoder_h_1, decoder_h_2, decoder_h_3, decoder_h_4 = decoder(code, decoder_h_1, decoder_h_2, decoder_h_3, decoder_h_4)
encoded2, encoder_h_1_2, encoder_h_2_2, encoder_h_3_2 = encoder2(res2, encoder_h_1_2, encoder_h_2_2, encoder_h_3_2)
code2 = binarizer2(encoded2)
output2, decoder_h_1_2, decoder_h_2_2, decoder_h_3_2, decoder_h_4_2 = decoder2(code2, decoder_h_1_2, decoder_h_2_2, decoder_h_3_2, decoder_h_4_2)
im += output.data.cpu()
res = res - output
codes.append(code.data.cpu().numpy())
im2 += output2.data.cpu()
res2 = res2 - output2
codes2.append(code2.data.cpu().numpy())
qbits+= (self.height/16*self.width/16 * 32)*self.batch_size
#print('code.shape',code.shape,'code2.shape', code2.shape)
for i in range(self.batch_size):
im_new = torch.cat((im, im2), 1)
#print(im_new.shape)
img = np.squeeze(im_new[i].numpy().clip(0, 1) * 255.0).astype(np.uint8).transpose(1, 2, 0)
list_patches_recons[i+batch*self.batch_size,:,:,:] = img
if self.op_bit_allocation and iters+1>=self.min_iters:
result =[]
for i in range(self.batch_size):
patch_test = np.squeeze(im[i].numpy().clip(0, 1) * 255.0).astype(np.uint8).transpose(1, 2, 0)
#patch_test=np.squeeze(im.numpy().clip(0, 1) * 255.0).astype(np.uint8).transpose(1, 2, 0)
patch_test = Image.fromarray(patch_test, mode = self.color_mode)
patch_test = patch_test.convert('RGB')
patch_ref = (patches[batch*batch_size+i]* 255.0).astype(np.uint8)
#patch_ref = Image.fromarray(patch_ref, mode='RGB')
patch_ref = Image.fromarray(patch_ref, mode = self.color_mode)
patch_ref = patch_ref.convert('RGB')
result.append(calc_metric(patch_ref,patch_test,'psnr'))
lower = min(result)
media = np.mean(result)
if media>=th and lower>=th*0.90:
print('iteração',iters+1,'PSNR médio', media)
break
bpp = qbits/(768*512)
codes1 = np.asarray(codes).flatten()
codes2 = np.asarray(codes2).flatten()
codes12 = np.concatenate((codes1, codes2), axis=0)
if batch==0:
code_img = codes12
code_img1 = codes1
code_img2 = codes2
else:
code_img = np.concatenate((code_img, codes12), axis=0)
code_img1 = np.concatenate((code_img1, codes1), axis=0)
code_img2 = np.concatenate((code_img2, codes2), axis=0)
code_img1 = list(map(int, code_img1))
latent1 = list(map(lambda x: x != -1, code_img1))
size_in1 = bitstring.Bits(latent1).length
compressed1 = gzip.compress(bitstring.Bits(latent1).bytes)
size_out1 = bitstring.Bits(compressed1).length
#print('entrada (bits)', size_in,'saída(bits)', size_out, 'Data compression ratio',size_in/size_out)
#bpp2 = size_out/(512*768)
code_img2 = list(map(int, code_img2))
latent2 = list(map(lambda x: x != -1, code_img2))
size_in2 = bitstring.Bits(latent2).length
compressed2 = gzip.compress(bitstring.Bits(latent2).bytes)
size_out2 = bitstring.Bits(compressed2).length
print('entrada (bits)', size_in1+size_in1,'saída(bits)', size_out1+size_out2, 'Data compression ratio',(size_in1+size_in2)/(size_out1 + size_out2))
bpp2 = (size_out1+size_out2)/(512*768)
return list_patches_recons,bpp,bpp2
