├── README.md ├── code └── lossy_multi_model │ ├── EDEH.py │ ├── basic_DL_op.py │ ├── creat_Long_context.py │ ├── entropy_codec.py │ ├── lifting97.py │ ├── main_testRGB.py │ ├── pixelcnn_2D.py │ ├── pixelcnn_2D_context.py │ ├── scale_quant.py │ └── wavelet.py ├── figs ├── Kodak.PNG └── Tecnick.PNG └── model └── README.md /README.md: -------------------------------------------------------------------------------- 1 | # Versatile-Image-Compression 2 | 3 | *NOTE: A new pytorch implementation of iWave++ has been released at (https://gitlab.com/iWave/iwave), including training and testing code, and pre-trained models. This pytorch implementation is also used as the reference software for IEEE 1857.11 Standard for Neural Network-Based Image Coding.* 4 | 5 | This repo provides the official implementation of "[End-to-End Optimized Versatile Image Compression With Wavelet-Like Transform](https://ieeexplore.ieee.org/document/9204799)". 6 | 7 | Accepted by IEEE TPAMI. 8 | 9 | Author: Haichuan Ma, Dong Liu, Ning Yan, Houqiang Li, Feng Wu 10 | 11 | ## **BibTeX** 12 | 13 | @article{ma2020end, 14 | title={End-to-End Optimized Versatile Image Compression With Wavelet-Like Transform}, 15 | author={Ma, Haichuan and Liu, Dong and Yan, Ning and Li, Houqiang and Wu, Feng}, 16 | journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 17 | year={2020}, 18 | publisher={IEEE} 19 | } 20 | 21 | ## **Update Notes** 22 | 23 | 2022.10.28 iWave++ has been accepted as one of the three reference softwares for IEEE 1857.11 Standard for Neural Network-Based Image Coding. Pytorch code (including both training and testing code) and pre-trained models can be found at (https://gitlab.com/iWave/iwave). To further improve compression performance, a series of improvements (such as an enhanced context model and an enhanced de-quant model) have been introduced in the new pytorch implementation compared to the original tensorflow version. 24 | 25 | 2020.9.4 Upload code and models of **Lossy multi-model iWave++**. 26 | 27 | 2020.8.26 Init this repo. 28 | 29 | ## **How To Test** 30 | 0. Dependencies. We test with MIT deepo docker image. 31 | 32 | 1. Clone this github repo. 33 | 34 | 2. Place Test images. (The code now only supports images whose border length is a multiple of 16. However, it is very simple to support arbitrary boundary lengths by padding.) 35 | 36 | 3. Download models. See **model** folder. 37 | 38 | 4. python main_testRGB.py. (The path in main_testRGB.py needs to be modified. Please refer to the code.) 39 | 40 | 41 | ## **Results** 42 | 43 | iWave++ outperforms [Joint](http://papers.nips.cc/paper/8275-joint-autoregressive-and-hierarchical-priors-for-learned-image-compression), [Variational](https://arxiv.org/abs/1802.01436), and [iWave](https://ieeexplore.ieee.org/abstract/document/8931632). For more information, please refer to the paper. 44 | 45 | 1. RGB PSNR on Kodak dataset. 46 | 47 | ![image](https://github.com/mahaichuan/Versatile-Image-Compression/blob/master/figs/Kodak.PNG) 48 | 49 | 2. RGB PSNR on Tecnick dataset. 50 | 51 | ![image](https://github.com/mahaichuan/Versatile-Image-Compression/blob/master/figs/Tecnick.PNG) 52 | -------------------------------------------------------------------------------- /code/lossy_multi_model/EDEH.py: -------------------------------------------------------------------------------- 1 | import tensorflow.contrib.slim as slim 2 | import tensorflow as tf 3 | import numpy as np 4 | import shutil 5 | 6 | # basic op of deep learning 7 | def weight_variable(name,shape): 8 | initial = tf.truncated_normal(shape, stddev=0.01) 9 | return tf.get_variable(name, initializer=initial) 10 | 11 | def bias_variable(name,shape): 12 | initial = tf.constant(0., shape=shape) 13 | return tf.get_variable(name, initializer=initial) 14 | 15 | def conv2d(x, W): 16 | return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME') 17 | 18 | # basic residual block of EDEH 19 | def resBlock(x, channels=16): 20 | 21 | w = weight_variable('conv1',[3,3,channels,channels]) 22 | b = bias_variable('bias1',[channels]) 23 | 24 | tmp = conv2d(x,w) + b 25 | 26 | tmp = tf.nn.relu(tmp) 27 | 28 | 29 | w = weight_variable('conv2',[3, 3, channels, channels]) 30 | b = bias_variable('bias2',[channels]) 31 | 32 | tmp = conv2d(tmp,w) + b 33 | 34 | return x + tmp 35 | 36 | # Network of EDEH 37 | 38 | def EDEH(x, num_layers=6,feature_size=64,output_channels=1): 39 | 40 | in_x = x 41 | 42 | # //////////////// conv: before residual blocks 43 | 44 | w = weight_variable('conv1',[3,3,1,feature_size]) 45 | b = bias_variable('bias1',[feature_size]) 46 | 47 | x = conv2d(x,w) + b 48 | conv_1 = x 49 | 50 | 51 | # ////////////// residual blocks 52 | 53 | for i in range(num_layers): 54 | with tf.variable_scope('resBlock_' + str(i)): 55 | x = resBlock(x, feature_size) 56 | 57 | # ///////////// One more convolution, and then we add the output of our first conv layer 58 | w = weight_variable('conv2',[3, 3, feature_size, feature_size]) 59 | b = bias_variable('bias2',[feature_size]) 60 | x = conv2d(x,w) + b 61 | x = x + conv_1 62 | 63 | # //////////// one more convolution after direct addition 64 | w = weight_variable('conv3',[3,3,feature_size,output_channels]) 65 | b = bias_variable('bias3',[output_channels]) 66 | x = conv2d(x,w) + b 67 | 68 | x = x + in_x 69 | 70 | return x -------------------------------------------------------------------------------- /code/lossy_multi_model/basic_DL_op.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | 11 | ################# we define the basic DL operations here 12 | 13 | def weight_variable(name,shape,std): 14 | 15 | initial = tf.truncated_normal(shape, stddev=std) 16 | 17 | return tf.get_variable(name, initializer=initial) 18 | 19 | 20 | def bias_variable(name, shape): 21 | 22 | initial = tf.constant(0., shape=shape) 23 | 24 | return tf.get_variable(name, initializer=initial) 25 | 26 | 27 | def conv2d(x, W, stride=1): 28 | 29 | return tf.nn.conv2d(x, W, strides=[1, stride, stride, 1], padding='SAME') 30 | 31 | 32 | def conv2d_pad(x, W, stride=1): 33 | 34 | paddings = tf.constant([[0,0],[1,1],[1,1],[0,0]]) 35 | 36 | x = tf.pad(x, paddings=paddings,mode="REFLECT") 37 | 38 | return tf.nn.conv2d(x, W, strides=[1, stride, stride, 1], padding='VALID') 39 | # 40 | # def conv3d(x,W): 41 | # 42 | # return tf.nn.conv3d(input=x, filter=W, strides=[1,1,1,1,1], padding='SAME') -------------------------------------------------------------------------------- /code/lossy_multi_model/creat_Long_context.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import basic_DL_op 11 | 12 | def lstm_logic(x, c): 13 | 14 | i = tf.sigmoid(x) 15 | f = tf.sigmoid(x) 16 | o = tf.sigmoid(x) 17 | g = tf.tanh(x) 18 | 19 | c = f*c+i*g 20 | 21 | h = o*tf.tanh(c) 22 | 23 | 24 | return c, h 25 | 26 | def lstm_layer(x, h, c, in_num, out_num): 27 | 28 | # the first layer: input 29 | 30 | w = basic_DL_op.weight_variable('conv1', [3, 3, in_num, out_num], 0.01) 31 | 32 | x = basic_DL_op.conv2d(x, w) 33 | 34 | # the first layer: state 35 | 36 | w = basic_DL_op.weight_variable('conv2', [3, 3, out_num, out_num], 0.01) 37 | 38 | h = basic_DL_op.conv2d(h, w) 39 | 40 | b = basic_DL_op.bias_variable('bias', [out_num]) 41 | 42 | 43 | 44 | c, h = lstm_logic(x + h + b, c) 45 | 46 | 47 | return c, h 48 | 49 | 50 | def context_single_band(x1, h1, c1, h2, c2, h3, c3, bit_map): 51 | 52 | with tf.variable_scope('LSTM_' + str(1)): 53 | c1, h1 = lstm_layer(x1, h1, c1, 1, int(32*bit_map)) 54 | 55 | with tf.variable_scope('LSTM_' + str(2)): 56 | c2, h2 = lstm_layer(h1, h2, c2, int(32*bit_map), int(32 * bit_map)) 57 | 58 | with tf.variable_scope('LSTM_' + str(3)): 59 | c3, h3 = lstm_layer(h2, h3, c3, int(32 * bit_map), 1) 60 | 61 | return h1, c1, h2, c2, h3, c3 62 | 63 | 64 | def context_single_band_reuse(x1, h1, c1, h2, c2, h3, c3, bit_map): 65 | 66 | with tf.variable_scope('LSTM_' + str(1), reuse=True): 67 | c1, h1 = lstm_layer(x1, h1, c1, 1, int(32*bit_map)) 68 | 69 | with tf.variable_scope('LSTM_' + str(2), reuse=True): 70 | c2, h2 = lstm_layer(h1, h2, c2, int(32*bit_map), int(32 * bit_map)) 71 | 72 | with tf.variable_scope('LSTM_' + str(3), reuse=True): 73 | c3, h3 = lstm_layer(h2, h3, c3, int(32 * bit_map), 1) 74 | 75 | return h1, c1, h2, c2, h3, c3 76 | 77 | def deconv_layer(x): 78 | 79 | x_shape = x.get_shape().as_list() 80 | 81 | kernel = basic_DL_op.weight_variable('deconv', [3, 3, x_shape[3], x_shape[3]], 0.01) 82 | 83 | x = tf.nn.conv2d_transpose(x, kernel, output_shape=[x_shape[0], int(x_shape[1]*2), int(x_shape[2]*2), x_shape[3]], strides=[1, 2, 2, 1], padding="SAME") 84 | 85 | return x 86 | 87 | 88 | 89 | def context_all(LL, HL_collection, LH_collection, HH_collection, bit_map, static_QP): 90 | 91 | c_HL = [] 92 | c_LH = [] 93 | c_HH = [] 94 | 95 | x_shape = LL.get_shape().as_list() 96 | 97 | h1 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], int(32 * bit_map)], dtype=tf.float32) 98 | c1 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], int(32 * bit_map)], dtype=tf.float32) 99 | 100 | h2 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], int(32 * bit_map)], dtype=tf.float32) 101 | c2 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], int(32 * bit_map)], dtype=tf.float32) 102 | 103 | h3 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], 1], dtype=tf.float32) 104 | c3 = tf.zeros(shape=[x_shape[0], x_shape[1], x_shape[2], 1], dtype=tf.float32) 105 | 106 | 107 | h1, c1, h2, c2, h3, c3 = context_single_band(LL / static_QP, h1, c1, h2, c2, h3, c3, bit_map) 108 | 109 | c_HL.append(h3) 110 | 111 | for j in range(4): 112 | 113 | i = 4 - 1 - j 114 | 115 | h1, c1, h2, c2, h3, c3 = context_single_band_reuse(HL_collection[i]/ static_QP, h1, c1, h2, c2, h3, c3, bit_map) 116 | 117 | c_LH.append(h3) 118 | 119 | h1, c1, h2, c2, h3, c3 = context_single_band_reuse(LH_collection[i]/ static_QP, h1, c1, h2, c2, h3, c3, bit_map) 120 | 121 | c_HH.append(h3) 122 | 123 | h1, c1, h2, c2, h3, c3 = context_single_band_reuse(HH_collection[i]/ static_QP, h1, c1, h2, c2, h3, c3, bit_map) 124 | 125 | with tf.variable_scope('Deconv_h1' + str(j)): 126 | h1 = deconv_layer(h1) 127 | with tf.variable_scope('Deconv_c1' + str(j)): 128 | c1 = deconv_layer(c1) 129 | with tf.variable_scope('Deconv_h2' + str(j)): 130 | h2 = deconv_layer(h2) 131 | with tf.variable_scope('Deconv_c2' + str(j)): 132 | c2 = deconv_layer(c2) 133 | with tf.variable_scope('Deconv_h3' + str(j)): 134 | h3 = deconv_layer(h3) 135 | with tf.variable_scope('Deconv_c3' + str(j)): 136 | c3 = deconv_layer(c3) 137 | 138 | c_HL.append(h3) 139 | 140 | 141 | return c_HL, c_LH, c_HH -------------------------------------------------------------------------------- /code/lossy_multi_model/entropy_codec.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import pixelcnn_2D 11 | import pixelcnn_2D_context 12 | 13 | 14 | def codec(LL, HL_collection, LH_collection, HH_collection, c_HL, c_LH, c_HH, static_QP): 15 | 16 | with tf.variable_scope('LL'): 17 | 18 | ce_loss = pixelcnn_2D.mask_2D_layer(LL, static_QP) 19 | 20 | for j in range(4): 21 | 22 | i = 4 - 1 - j 23 | 24 | c = tf.pow(tf.pow(2,(4-1-i)), 2) 25 | 26 | c = tf.cast(c, dtype=tf.float32) 27 | 28 | with tf.variable_scope('HL'+str(i)): 29 | ce_loss = pixelcnn_2D_context.mask_2D_layer(HL_collection[i], static_QP, c_HL[j]) * c + ce_loss 30 | 31 | with tf.variable_scope('LH'+str(i)): 32 | ce_loss = pixelcnn_2D_context.mask_2D_layer(LH_collection[i], static_QP, c_LH[j]) * c + ce_loss 33 | 34 | with tf.variable_scope('HH'+str(i)): 35 | ce_loss = pixelcnn_2D_context.mask_2D_layer(HH_collection[i], static_QP, c_HH[j]) * c + ce_loss 36 | 37 | return ce_loss / 256. -------------------------------------------------------------------------------- /code/lossy_multi_model/lifting97.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import basic_DL_op 11 | 12 | ######################## lifting 97 forward and inverse transform 13 | 14 | lifting_coeff = [-1.586134342059924, -0.052980118572961, 0.882911075530934, 0.443506852043971, 0.869864451624781, 1.149604398860241] # bior4.4 15 | # lifting_coeff = [-12993./8193., -434./8193., 7233./8193., 3633./8193., 5038./8193., 6659./8193.] # OPENJPEG 16 | # lifting_coeff = [-1.586134342059924, -0.052980118572961, 0.882911075530934, 0.443506852043971, 1.230174104914001, 0.812893066115961] # CDF9/7 17 | 18 | trainable_set = True 19 | 20 | def p_block(x): 21 | 22 | w = basic_DL_op.weight_variable('conv1',[3,3,1,16],0.01) 23 | b = basic_DL_op.bias_variable('bias1', [16]) 24 | tmp = basic_DL_op.conv2d_pad(x,w) + b 25 | 26 | conv1 = tmp 27 | 28 | tmp = tf.nn.tanh(tmp) 29 | w = basic_DL_op.weight_variable('conv2', [3, 3, 16, 16], 0.01) 30 | b = basic_DL_op.bias_variable('bias2', [16]) 31 | tmp = basic_DL_op.conv2d_pad(tmp, w) + b 32 | 33 | tmp = tf.nn.tanh(tmp) 34 | w = basic_DL_op.weight_variable('conv3', [3, 3, 16, 16], 0.01) 35 | b = basic_DL_op.bias_variable('bias3', [16]) 36 | tmp = basic_DL_op.conv2d_pad(tmp, w) + b 37 | 38 | tmp = conv1 + tmp 39 | 40 | w = basic_DL_op.weight_variable('conv4', [3, 3, 16, 1], 0.01) 41 | b = basic_DL_op.bias_variable('bias4', [1]) 42 | out = basic_DL_op.conv2d_pad(tmp, w) + b 43 | 44 | return out 45 | 46 | def lifting97_forward(L, H): 47 | 48 | # first lifting step 49 | 50 | # with tf.variable_scope('p_1'): 51 | # 52 | # L_net = p_block(L/256.)*256. 53 | 54 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 55 | 56 | tmp = tf.pad(L, paddings=paddings, mode="REFLECT") 57 | 58 | initial = [[0.0], 59 | [lifting_coeff[0]], 60 | [lifting_coeff[0]]] 61 | 62 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 63 | 64 | w = tf.get_variable('cdf97_0', initializer=initial, trainable=trainable_set) 65 | 66 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 67 | 68 | with tf.variable_scope('p_1'): 69 | 70 | L_net = p_block(skip/256.)*256. 71 | 72 | H = H + skip + L_net * 0.1 73 | 74 | 75 | 76 | # with tf.variable_scope('u_1'): 77 | # 78 | # H_net = p_block(H/256.)*256. 79 | 80 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 81 | 82 | tmp = tf.pad(H, paddings=paddings, mode="REFLECT") 83 | 84 | initial = [[lifting_coeff[1]], 85 | [lifting_coeff[1]], 86 | [0.0]] 87 | 88 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 89 | 90 | w = tf.get_variable('cdf97_1', initializer=initial, trainable=trainable_set) 91 | 92 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 93 | 94 | with tf.variable_scope('u_1'): 95 | 96 | H_net = p_block(skip/256.)*256. 97 | 98 | L = L + skip + H_net * 0.1 99 | 100 | # second lifting step 101 | 102 | # with tf.variable_scope('p_2'): 103 | # 104 | # L_net = p_block(L/256.)*256. 105 | 106 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 107 | 108 | tmp = tf.pad(L, paddings=paddings, mode="REFLECT") 109 | 110 | initial = [[0.0], 111 | [lifting_coeff[2]], 112 | [lifting_coeff[2]]] 113 | 114 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 115 | 116 | w = tf.get_variable('cdf97_2', initializer=initial, trainable=trainable_set) 117 | 118 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 119 | 120 | with tf.variable_scope('p_2'): 121 | 122 | L_net = p_block(skip/256.)*256. 123 | 124 | H = H + skip + L_net * 0.1 125 | 126 | 127 | 128 | # with tf.variable_scope('u_2'): 129 | # 130 | # H_net = p_block(H/256.)*256. 131 | 132 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 133 | 134 | tmp = tf.pad(H, paddings=paddings, mode="REFLECT") 135 | 136 | initial = [[lifting_coeff[3]], 137 | [lifting_coeff[3]], 138 | [0.0]] 139 | 140 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 141 | 142 | w = tf.get_variable('cdf97_3', initializer=initial, trainable=trainable_set) 143 | 144 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 145 | 146 | with tf.variable_scope('u_2'): 147 | 148 | H_net = p_block(skip/256.)*256. 149 | 150 | L = L + skip + H_net * 0.1 151 | 152 | # scaling step 153 | 154 | n_h = tf.get_variable('n_h', initializer=0.0, trainable=trainable_set) 155 | 156 | n_h = lifting_coeff[4] + n_h * 0.1 157 | 158 | H = H * n_h 159 | 160 | n_l = tf.get_variable('n_l', initializer=0.0, trainable=trainable_set) 161 | 162 | n_l = lifting_coeff[5] + n_l * 0.1 163 | 164 | L = L * n_l 165 | 166 | return L, H 167 | 168 | 169 | def lifting97_inverse(L, H): 170 | 171 | # scaling step 172 | 173 | n_h = tf.get_variable('n_h', initializer=0.0, trainable=trainable_set) 174 | 175 | n_h = lifting_coeff[4] + n_h * 0.1 176 | 177 | H = H / n_h 178 | 179 | n_l = tf.get_variable('n_l', initializer=0.0, trainable=trainable_set) 180 | 181 | n_l = lifting_coeff[5] + n_l * 0.1 182 | 183 | L = L / n_l 184 | 185 | # second lifting step 186 | 187 | # with tf.variable_scope('u_2'): 188 | # H_net = p_block(H / 256.) * 256. 189 | 190 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 191 | 192 | tmp = tf.pad(H, paddings=paddings, mode="REFLECT") 193 | 194 | initial = [[lifting_coeff[3]], 195 | [lifting_coeff[3]], 196 | [0.0]] 197 | 198 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 199 | 200 | w = tf.get_variable('cdf97_3', initializer=initial, trainable=trainable_set) 201 | 202 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 203 | 204 | with tf.variable_scope('u_2'): 205 | H_net = p_block(skip / 256.) * 256. 206 | 207 | L = L - skip - H_net * 0.1 208 | 209 | 210 | 211 | # with tf.variable_scope('p_2'): 212 | # L_net = p_block(L / 256.) * 256. 213 | 214 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 215 | 216 | tmp = tf.pad(L, paddings=paddings, mode="REFLECT") 217 | 218 | initial = [[0.0], 219 | [lifting_coeff[2]], 220 | [lifting_coeff[2]]] 221 | 222 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 223 | 224 | w = tf.get_variable('cdf97_2', initializer=initial, trainable=trainable_set) 225 | 226 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 227 | 228 | with tf.variable_scope('p_2'): 229 | L_net = p_block(skip / 256.) * 256. 230 | 231 | H = H - skip - L_net * 0.1 232 | 233 | # first lifting step 234 | 235 | # with tf.variable_scope('u_1'): 236 | # H_net = p_block(H / 256.) * 256. 237 | 238 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 239 | 240 | tmp = tf.pad(H, paddings=paddings, mode="REFLECT") 241 | 242 | initial = [[lifting_coeff[1]], 243 | [lifting_coeff[1]], 244 | [0.0]] 245 | 246 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 247 | 248 | w = tf.get_variable('cdf97_1', initializer=initial, trainable=trainable_set) 249 | 250 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 251 | 252 | with tf.variable_scope('u_1'): 253 | H_net = p_block(skip / 256.) * 256. 254 | 255 | L = L - skip - H_net * 0.1 256 | 257 | 258 | # with tf.variable_scope('p_1'): 259 | # L_net = p_block(L / 256.) * 256. 260 | 261 | paddings = tf.constant([[0, 0], [1, 1], [0, 0], [0, 0]]) 262 | 263 | tmp = tf.pad(L, paddings=paddings, mode="REFLECT") 264 | 265 | initial = [[0.0], 266 | [lifting_coeff[0]], 267 | [lifting_coeff[0]]] 268 | 269 | initial = tf.reshape(initial, shape=[3, 1, 1, 1]) 270 | 271 | w = tf.get_variable('cdf97_0', initializer=initial, trainable=trainable_set) 272 | 273 | skip = tf.nn.conv2d(tmp, w, strides=[1, 1, 1, 1], padding='VALID') 274 | 275 | with tf.variable_scope('p_1'): 276 | L_net = p_block(skip / 256.) * 256. 277 | 278 | H = H - skip - L_net * 0.1 279 | 280 | return L, H 281 | -------------------------------------------------------------------------------- /code/lossy_multi_model/main_testRGB.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | 11 | import wavelet 12 | import scale_quant 13 | import entropy_codec 14 | import creat_Long_context 15 | import EDEH 16 | 17 | decomposition_step = 4 18 | 19 | def graph(x): 20 | 21 | LL = x 22 | 23 | # Convert RGB to YUV 24 | 25 | convert_mat = np.array([[0.299, 0.587, 0.114], [-0.169, -0.331, 0.499], [0.499, -0.418, -0.0813]], dtype=np.float32) 26 | 27 | inverse_convert_mat = np.array([[1.0, 0.0, 1.402], [1.0, -0.344, -0.714], [1.0, 1.772, 0.0]], dtype=np.float32) 28 | 29 | Y = LL[:,:,:,0:1] * convert_mat[0,0] + LL[:,:,:,1:2] * convert_mat[0,1] + LL[:,:,:,2:3] * convert_mat[0,2] 30 | U = LL[:, :, :, 0:1] * convert_mat[1, 0] + LL[:, :, :, 1:2] * convert_mat[1, 1] + LL[:, :, :, 2:3] * convert_mat[1, 2] + 128. 31 | V = LL[:, :, :, 0:1] * convert_mat[2, 0] + LL[:, :, :, 1:2] * convert_mat[2, 1] + LL[:, :, :, 2:3] * convert_mat[2, 2] + 128. 32 | 33 | # ########################### for Y 34 | 35 | LL = Y 36 | 37 | HL_collection = [] 38 | LH_collection = [] 39 | HH_collection = [] 40 | 41 | # forward transform, bior4.4 42 | 43 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 44 | 45 | for i in range(decomposition_step): 46 | LL, HL, LH, HH = wavelet.decomposition(LL) 47 | 48 | HL_collection.append(HL) 49 | LH_collection.append(LH) 50 | HH_collection.append(HH) 51 | 52 | # quant, [x*QP] 53 | 54 | with tf.variable_scope('quant', reuse=tf.AUTO_REUSE): 55 | 56 | static_QP = tf.get_variable('static_QP', initializer=1 / 16) 57 | 58 | # static_QP = 1 / 39.2063 59 | 60 | LL, HL_collection, LH_collection, HH_collection = scale_quant.quant(LL, HL_collection, LH_collection, HH_collection, 61 | static_QP) 62 | 63 | with tf.variable_scope('long_context', reuse=tf.AUTO_REUSE): 64 | 65 | c_HL, c_LH, c_HH = creat_Long_context.context_all(LL, HL_collection, LH_collection, HH_collection, bit_map=1, 66 | static_QP=static_QP * 4096.) 67 | 68 | with tf.variable_scope('ce_loss', reuse=tf.AUTO_REUSE): 69 | 70 | ce_loss_Y = entropy_codec.codec(LL, HL_collection, LH_collection, HH_collection, c_HL, c_LH, c_HH, 71 | static_QP=static_QP * 4096.) 72 | 73 | LL, HL_collection, LH_collection, HH_collection = scale_quant.de_quant(LL, HL_collection, LH_collection, 74 | HH_collection, static_QP) 75 | 76 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 77 | 78 | for j in range(decomposition_step): 79 | i = decomposition_step - 1 - j 80 | 81 | LL = wavelet.reconstruct(LL, HL_collection[i], LH_collection[i], HH_collection[i]) 82 | 83 | LL = LL / 255. - 0.5 84 | 85 | with tf.variable_scope('post_process', reuse=tf.AUTO_REUSE): 86 | 87 | LL = EDEH.EDEH(LL) 88 | 89 | LL = (LL + 0.5) * 255. 90 | 91 | Y = LL 92 | 93 | # ########################## for U 94 | 95 | LL = U 96 | 97 | HL_collection = [] 98 | LH_collection = [] 99 | HH_collection = [] 100 | 101 | # forward transform, bior4.4 102 | 103 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 104 | 105 | for i in range(decomposition_step): 106 | LL, HL, LH, HH = wavelet.decomposition(LL) 107 | 108 | HL_collection.append(HL) 109 | LH_collection.append(LH) 110 | HH_collection.append(HH) 111 | 112 | # quant, [x*QP] 113 | 114 | with tf.variable_scope('quant', reuse=tf.AUTO_REUSE): 115 | 116 | static_QP = tf.get_variable('static_QP', initializer=1 / 16) 117 | 118 | # static_QP = 1 / 39.2063 119 | 120 | LL, HL_collection, LH_collection, HH_collection = scale_quant.quant(LL, HL_collection, LH_collection, HH_collection, 121 | static_QP) 122 | 123 | with tf.variable_scope('long_context', reuse=tf.AUTO_REUSE): 124 | 125 | c_HL, c_LH, c_HH = creat_Long_context.context_all(LL, HL_collection, LH_collection, HH_collection, bit_map=1, 126 | static_QP=static_QP * 4096.) 127 | 128 | with tf.variable_scope('ce_loss', reuse=tf.AUTO_REUSE): 129 | 130 | ce_loss_U = entropy_codec.codec(LL, HL_collection, LH_collection, HH_collection, c_HL, c_LH, c_HH, 131 | static_QP=static_QP * 4096.) 132 | 133 | LL, HL_collection, LH_collection, HH_collection = scale_quant.de_quant(LL, HL_collection, LH_collection, 134 | HH_collection, static_QP) 135 | 136 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 137 | 138 | for j in range(decomposition_step): 139 | i = decomposition_step - 1 - j 140 | 141 | LL = wavelet.reconstruct(LL, HL_collection[i], LH_collection[i], HH_collection[i]) 142 | 143 | LL = LL / 255. - 0.5 144 | 145 | with tf.variable_scope('post_process', reuse=tf.AUTO_REUSE): 146 | 147 | LL = EDEH.EDEH(LL) 148 | 149 | LL = (LL + 0.5) * 255. 150 | 151 | U = LL 152 | 153 | # ########################## for V 154 | 155 | LL = V 156 | 157 | HL_collection = [] 158 | LH_collection = [] 159 | HH_collection = [] 160 | 161 | # forward transform, bior4.4 162 | 163 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 164 | 165 | for i in range(decomposition_step): 166 | LL, HL, LH, HH = wavelet.decomposition(LL) 167 | 168 | HL_collection.append(HL) 169 | LH_collection.append(LH) 170 | HH_collection.append(HH) 171 | 172 | # quant, [x*QP] 173 | 174 | with tf.variable_scope('quant', reuse=tf.AUTO_REUSE): 175 | 176 | static_QP = tf.get_variable('static_QP', initializer=1 / 16) 177 | 178 | # static_QP = 1 / 39.2063 179 | 180 | LL, HL_collection, LH_collection, HH_collection = scale_quant.quant(LL, HL_collection, LH_collection, HH_collection, 181 | static_QP) 182 | 183 | with tf.variable_scope('long_context', reuse=tf.AUTO_REUSE): 184 | 185 | c_HL, c_LH, c_HH = creat_Long_context.context_all(LL, HL_collection, LH_collection, HH_collection, bit_map=1, 186 | static_QP=static_QP * 4096.) 187 | 188 | with tf.variable_scope('ce_loss', reuse=tf.AUTO_REUSE): 189 | 190 | ce_loss_V = entropy_codec.codec(LL, HL_collection, LH_collection, HH_collection, c_HL, c_LH, c_HH, 191 | static_QP=static_QP * 4096.) 192 | 193 | LL, HL_collection, LH_collection, HH_collection = scale_quant.de_quant(LL, HL_collection, LH_collection, 194 | HH_collection, static_QP) 195 | 196 | with tf.variable_scope('wavelet', reuse=tf.AUTO_REUSE): 197 | 198 | for j in range(decomposition_step): 199 | i = decomposition_step - 1 - j 200 | 201 | LL = wavelet.reconstruct(LL, HL_collection[i], LH_collection[i], HH_collection[i]) 202 | 203 | LL = LL / 255. - 0.5 204 | 205 | with tf.variable_scope('post_process', reuse=tf.AUTO_REUSE): 206 | 207 | LL = EDEH.EDEH(LL) 208 | 209 | LL = (LL + 0.5) * 255. 210 | 211 | V = LL 212 | 213 | 214 | # Convert YUV to RGB 215 | 216 | R = Y * inverse_convert_mat[0, 0] + (U - 128.) * inverse_convert_mat[0, 1] + (V - 128.) * inverse_convert_mat[0, 2] 217 | G = Y * inverse_convert_mat[1, 0] + (U - 128.) * inverse_convert_mat[1, 1] + (V - 128.) * inverse_convert_mat[1, 2] 218 | B = Y * inverse_convert_mat[2, 0] + (U - 128.) * inverse_convert_mat[2, 1] + (V - 128.) * inverse_convert_mat[2, 2] 219 | 220 | LL = tf.concat([R, G, B], axis=3) 221 | 222 | LL = tf.clip_by_value(LL, 0., 255.) 223 | 224 | 225 | 226 | d_loss = tf.losses.mean_squared_error(x, LL) 227 | 228 | ce_loss = ce_loss_Y + ce_loss_U + ce_loss_V 229 | 230 | return ce_loss/tf.log(2.), d_loss, ce_loss_Y/tf.log(2.), ce_loss_U/tf.log(2.), ce_loss_V/tf.log(2.), LL 231 | 232 | 233 | 234 | h_in = 1200 235 | w_in = 1200 236 | 237 | x = tf.placeholder(tf.float32, [1, h_in, w_in, 3]) 238 | 239 | ce_loss, d_loss, ce_loss_Y, ce_loss_U, ce_loss_V, recon = graph(x) 240 | 241 | saver = tf.train.Saver() 242 | 243 | rate_all = [] 244 | psnr_all = [] 245 | 246 | with tf.Session() as sess: 247 | 248 | saver.restore(sess, 'your-model-path') 249 | 250 | for batch_index in range(100): 251 | 252 | i = batch_index + 1 253 | 254 | print('img_ID:', i) 255 | 256 | img = Image.open('Tecnick-path/RGB_OR_1200x1200_'+str(i).zfill(3)+'.png') 257 | 258 | img = np.asarray(img, dtype=np.float32) 259 | 260 | img = np.reshape(img, (1, h_in, w_in, 3)) 261 | 262 | rate, d_eval, r_Y, r_U, r_V, recon_eval = sess.run([ce_loss, d_loss, ce_loss_Y, ce_loss_U, ce_loss_V, recon], feed_dict={x: img}) 263 | 264 | # recon_eval = recon_eval[0, :, :, :] 265 | # 266 | # img = Image.fromarray(np.uint8(recon_eval)) 267 | # 268 | # img.save('./recon/'+str(i)+'.png') 269 | 270 | psnr_ = 10 * log(255.*255. / d_eval, 10) 271 | 272 | print('PSNR:', psnr_) 273 | print('rate:', rate) 274 | 275 | rate_all.append(rate) 276 | psnr_all.append(psnr_) 277 | 278 | rate_all = np.array(rate_all) 279 | psnr_all = np.array(psnr_all) 280 | print('rate_mean:', np.mean(rate_all)) 281 | print('psnr_mean:', np.mean(psnr_all)) 282 | 283 | sess.close() 284 | -------------------------------------------------------------------------------- /code/lossy_multi_model/pixelcnn_2D.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import basic_DL_op 11 | 12 | def mask_2D_resiBlock(x, filter_nums): 13 | 14 | w = basic_DL_op.weight_variable('conv1', [3, 3, filter_nums, filter_nums], 0.01) 15 | 16 | 17 | mask = [[1,1,1], 18 | [1,1,0], 19 | [0,0,0]] 20 | 21 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 22 | 23 | mask = tf.tile(mask, multiples=[1, 1, filter_nums, filter_nums]) 24 | 25 | mask = tf.cast(mask, dtype=tf.float32) 26 | 27 | w = w * mask 28 | 29 | b = basic_DL_op.bias_variable('bias1', [filter_nums]) 30 | 31 | c = basic_DL_op.conv2d(x, w) + b 32 | 33 | c = tf.nn.relu(c) 34 | 35 | 36 | 37 | w = basic_DL_op.weight_variable('conv2', [3, 3, filter_nums, filter_nums], 0.01) 38 | 39 | w = w * mask 40 | 41 | b = basic_DL_op.bias_variable('bias2', [filter_nums]) 42 | 43 | c = basic_DL_op.conv2d(c, w) + b 44 | 45 | return x + c 46 | 47 | def cal_cdf(logits, h, b, a): 48 | 49 | shape = logits.get_shape().as_list() 50 | logits = tf.reshape(logits, [shape[0], shape[1], shape[2], shape[3], 1]) 51 | 52 | logits = tf.matmul(tf.reshape(h[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]), logits) 53 | logits = logits + tf.reshape(b[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]) 54 | logits = logits + tf.reshape(a[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 55 | 56 | logits = tf.matmul(tf.reshape(h[:,:,:,3:12], [shape[0], shape[1], shape[2], 3,3]), logits) 57 | logits = logits + tf.reshape(b[:,:,:,3:6], [shape[0], shape[1], shape[2], 3,1]) 58 | logits = logits + tf.reshape(a[:,:,:,3:6], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 59 | 60 | logits = tf.matmul(tf.reshape(h[:,:,:,12:21], [shape[0], shape[1], shape[2], 3,3]), logits) 61 | logits = logits + tf.reshape(b[:,:,:,6:9], [shape[0], shape[1], shape[2], 3,1]) 62 | logits = logits + tf.reshape(a[:,:,:,6:9], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 63 | 64 | logits = tf.matmul(tf.reshape(h[:,:,:,21:30], [shape[0], shape[1], shape[2], 3,3]), logits) 65 | logits = logits + tf.reshape(b[:,:,:,9:12], [shape[0], shape[1], shape[2], 3,1]) 66 | logits = logits + tf.reshape(a[:,:,:,9:12], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 67 | 68 | logits = tf.matmul(tf.reshape(h[:,:,:,30:33], [shape[0], shape[1], shape[2], 1,3]), logits) 69 | logits = logits + tf.reshape(b[:,:,:,12:13], [shape[0], shape[1], shape[2], 1,1]) 70 | 71 | logits = tf.sigmoid(logits) 72 | logits = tf.reshape(logits, [shape[0], shape[1], shape[2], shape[3]]) 73 | 74 | return logits 75 | 76 | def mask_2D_layer(x, static_QP, features = 128, resi_num = 2, para_num = 58): 77 | 78 | x = x / static_QP 79 | label = x 80 | 81 | # x = tf.stop_gradient(x) 82 | 83 | ################## layer 1, linear 84 | 85 | w = basic_DL_op.weight_variable('conv1', [3, 3, 1, features], 0.01) 86 | 87 | mask = [[1, 1, 1], 88 | [1, 0, 0], 89 | [0, 0, 0]] 90 | 91 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 92 | 93 | mask = tf.tile(mask, multiples=[1, 1, 1, features]) 94 | 95 | mask = tf.cast(mask, dtype=tf.float32) 96 | 97 | w = w * mask 98 | 99 | b = basic_DL_op.bias_variable('bias1', [features]) 100 | 101 | x = basic_DL_op.conv2d(x, w) + b 102 | 103 | 104 | conv1 = x 105 | 106 | ################## layers: resi_num resi_block 107 | 108 | for i in range(resi_num): 109 | with tf.variable_scope('resi_block' + str(i)): 110 | 111 | x = mask_2D_resiBlock(x, features) 112 | 113 | x = conv1 + x 114 | 115 | ################# conv: after skip connection, relu 116 | 117 | w = basic_DL_op.weight_variable('conv2', [3, 3, features, features], 0.01) 118 | 119 | mask = [[1, 1, 1], 120 | [1, 1, 0], 121 | [0, 0, 0]] 122 | 123 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 124 | 125 | mask = tf.tile(mask, multiples=[1, 1, features, features]) 126 | 127 | mask = tf.cast(mask, dtype=tf.float32) 128 | 129 | w = w * mask 130 | 131 | b = basic_DL_op.bias_variable('bias2', [features]) 132 | 133 | x = basic_DL_op.conv2d(x, w) + b 134 | 135 | x = tf.nn.relu(x) 136 | 137 | ################# convs: 1x1, relu/linear 138 | 139 | w = basic_DL_op.weight_variable('conv3', [1, 1, features, features], 0.01) 140 | 141 | b = basic_DL_op.bias_variable('bias3', [features]) 142 | 143 | x = basic_DL_op.conv2d(x, w) + b 144 | 145 | x = tf.nn.relu(x) 146 | 147 | 148 | w = basic_DL_op.weight_variable('conv4', [1, 1, features, features], 0.01) 149 | 150 | b = basic_DL_op.bias_variable('bias4', [features]) 151 | 152 | x = basic_DL_op.conv2d(x, w) + b 153 | 154 | x = tf.nn.relu(x) 155 | 156 | 157 | w = basic_DL_op.weight_variable('conv5', [1, 1, features, para_num], 0.01) 158 | 159 | b = basic_DL_op.bias_variable('bias5', [para_num]) 160 | 161 | x = basic_DL_op.conv2d(x, w) + b 162 | 163 | ################# cal the cdf with the output params 164 | 165 | h = tf.nn.softplus( x[:,:,:,0:33] ) 166 | b = x[:,:,:,33:46] 167 | a = tf.tanh( x[:,:,:,46:58] ) 168 | 169 | lower = label - 0.5 / static_QP 170 | high = label + 0.5 / static_QP 171 | 172 | lower = cal_cdf(lower, h, b, a) 173 | high = cal_cdf(high, h, b, a) 174 | 175 | prob = tf.maximum( (high - lower) , 1e-9) 176 | 177 | cross_entropy = -tf.reduce_mean(tf.log(prob)) 178 | 179 | return cross_entropy 180 | -------------------------------------------------------------------------------- /code/lossy_multi_model/pixelcnn_2D_context.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import basic_DL_op 11 | 12 | # we define the rate estimate functions here... 13 | 14 | def mask_2D_resiBlock(x, filter_nums): 15 | 16 | w = basic_DL_op.weight_variable('conv1', [3, 3, filter_nums, filter_nums], 0.01) 17 | 18 | mask = [[1,1,1], 19 | [1,1,0], 20 | [0,0,0]] 21 | 22 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 23 | 24 | mask = tf.tile(mask, multiples=[1, 1, filter_nums, filter_nums]) 25 | 26 | mask = tf.cast(mask, dtype=tf.float32) 27 | 28 | w = w * mask 29 | 30 | b = basic_DL_op.bias_variable('bias1', [filter_nums]) 31 | 32 | c = basic_DL_op.conv2d(x, w) + b 33 | 34 | c = tf.nn.relu(c) 35 | 36 | 37 | 38 | w = basic_DL_op.weight_variable('conv2', [3, 3, filter_nums, filter_nums], 0.01) 39 | 40 | w = w * mask 41 | 42 | b = basic_DL_op.bias_variable('bias2', [filter_nums]) 43 | 44 | c = basic_DL_op.conv2d(c, w) + b 45 | 46 | return x + c 47 | 48 | def resiBlock_2D_context(x, filter_nums): 49 | 50 | w = basic_DL_op.weight_variable('conv1_c', [3, 3, filter_nums, filter_nums], 0.01) 51 | 52 | b = basic_DL_op.bias_variable('bias1_c', [filter_nums]) 53 | 54 | c = basic_DL_op.conv2d(x, w) + b 55 | 56 | c = tf.nn.relu(c) 57 | 58 | w = basic_DL_op.weight_variable('conv2_c', [3, 3, filter_nums, filter_nums], 0.01) 59 | 60 | b = basic_DL_op.bias_variable('bias2_c', [filter_nums]) 61 | 62 | c = basic_DL_op.conv2d(c, w) + b 63 | 64 | return x + c 65 | 66 | def cal_cdf(logits, h, b, a): 67 | 68 | shape = logits.get_shape().as_list() 69 | logits = tf.reshape(logits, [shape[0], shape[1], shape[2], shape[3], 1]) 70 | 71 | logits = tf.matmul(tf.reshape(h[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]), logits) 72 | logits = logits + tf.reshape(b[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]) 73 | logits = logits + tf.reshape(a[:,:,:,0:3], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 74 | 75 | logits = tf.matmul(tf.reshape(h[:,:,:,3:12], [shape[0], shape[1], shape[2], 3,3]), logits) 76 | logits = logits + tf.reshape(b[:,:,:,3:6], [shape[0], shape[1], shape[2], 3,1]) 77 | logits = logits + tf.reshape(a[:,:,:,3:6], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 78 | 79 | logits = tf.matmul(tf.reshape(h[:,:,:,12:21], [shape[0], shape[1], shape[2], 3,3]), logits) 80 | logits = logits + tf.reshape(b[:,:,:,6:9], [shape[0], shape[1], shape[2], 3,1]) 81 | logits = logits + tf.reshape(a[:,:,:,6:9], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 82 | 83 | logits = tf.matmul(tf.reshape(h[:,:,:,21:30], [shape[0], shape[1], shape[2], 3,3]), logits) 84 | logits = logits + tf.reshape(b[:,:,:,9:12], [shape[0], shape[1], shape[2], 3,1]) 85 | logits = logits + tf.reshape(a[:,:,:,9:12], [shape[0], shape[1], shape[2], 3,1]) * tf.tanh(logits) 86 | 87 | logits = tf.matmul(tf.reshape(h[:,:,:,30:33], [shape[0], shape[1], shape[2], 1,3]), logits) 88 | logits = logits + tf.reshape(b[:,:,:,12:13], [shape[0], shape[1], shape[2], 1,1]) 89 | 90 | logits = tf.sigmoid(logits) 91 | logits = tf.reshape(logits, [shape[0], shape[1], shape[2], shape[3]]) 92 | 93 | return logits 94 | 95 | def mask_2D_layer(x, static_QP, context, features = 128, resi_num = 2, para_num = 58): 96 | 97 | x = x / static_QP 98 | label = x 99 | 100 | # x = tf.stop_gradient(x) 101 | 102 | ################## layer 1, linear 103 | 104 | w = basic_DL_op.weight_variable('conv1', [3, 3, 1, features], 0.01) 105 | w_c = basic_DL_op.weight_variable('conv1_c', [3, 3, 1, features], 0.01) 106 | 107 | mask = [[1, 1, 1], 108 | [1, 0, 0], 109 | [0, 0, 0]] 110 | 111 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 112 | 113 | mask = tf.tile(mask, multiples=[1, 1, 1, features]) 114 | 115 | mask = tf.cast(mask, dtype=tf.float32) 116 | 117 | w = w * mask 118 | 119 | b = basic_DL_op.bias_variable('bias1', [features]) 120 | b_c = basic_DL_op.bias_variable('bias1_c', [features]) 121 | 122 | x = basic_DL_op.conv2d(x, w) + b 123 | context = basic_DL_op.conv2d(context, w_c) + b_c 124 | 125 | 126 | conv1 = x 127 | x = x + context 128 | 129 | ################## layers: resi_num resi_block 130 | 131 | for i in range(resi_num): 132 | with tf.variable_scope('resi_block' + str(i)): 133 | 134 | x = mask_2D_resiBlock(x, features) 135 | 136 | context = resiBlock_2D_context(context, features) 137 | 138 | x = x + context 139 | 140 | x = conv1 + x 141 | 142 | ################# conv: after skip connection, relu 143 | 144 | w = basic_DL_op.weight_variable('conv2', [3, 3, features, features], 0.01) 145 | 146 | mask = [[1, 1, 1], 147 | [1, 1, 0], 148 | [0, 0, 0]] 149 | 150 | mask = tf.reshape(mask, shape=[3, 3, 1, 1]) 151 | 152 | mask = tf.tile(mask, multiples=[1, 1, features, features]) 153 | 154 | mask = tf.cast(mask, dtype=tf.float32) 155 | 156 | w = w * mask 157 | 158 | b = basic_DL_op.bias_variable('bias2', [features]) 159 | 160 | x = basic_DL_op.conv2d(x, w) + b 161 | 162 | x = tf.nn.relu(x) 163 | 164 | ################# convs: 1x1, relu/linear 165 | 166 | w = basic_DL_op.weight_variable('conv3', [1, 1, features, features], 0.01) 167 | 168 | b = basic_DL_op.bias_variable('bias3', [features]) 169 | 170 | x = basic_DL_op.conv2d(x, w) + b 171 | 172 | x = tf.nn.relu(x) 173 | 174 | 175 | w = basic_DL_op.weight_variable('conv4', [1, 1, features, features], 0.01) 176 | 177 | b = basic_DL_op.bias_variable('bias4', [features]) 178 | 179 | x = basic_DL_op.conv2d(x, w) + b 180 | 181 | x = tf.nn.relu(x) 182 | 183 | 184 | w = basic_DL_op.weight_variable('conv5', [1, 1, features, para_num], 0.01) 185 | 186 | b = basic_DL_op.bias_variable('bias5', [para_num]) 187 | 188 | x = basic_DL_op.conv2d(x, w) + b 189 | 190 | ################# cal the cdf with the output params 191 | 192 | h = tf.nn.softplus(x[:, :, :, 0:33]) 193 | b = x[:, :, :, 33:46] 194 | a = tf.tanh(x[:, :, :, 46:58]) 195 | 196 | lower = label - 0.5 / static_QP 197 | high = label + 0.5 / static_QP 198 | 199 | lower = cal_cdf(lower, h, b, a) 200 | high = cal_cdf(high, h, b, a) 201 | 202 | prob = tf.maximum((high - lower), 1e-9) 203 | 204 | cross_entropy = -tf.reduce_mean(tf.log(prob)) 205 | 206 | return cross_entropy 207 | -------------------------------------------------------------------------------- /code/lossy_multi_model/scale_quant.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import basic_DL_op 11 | 12 | clip_value = 8192. 13 | 14 | def quant(LL, HL_collection, LH_collection, HH_collection, s): 15 | 16 | with tf.get_default_graph().gradient_override_map({"Round": "Identity"}): 17 | 18 | HL_collection[0] = tf.round( s * tf.clip_by_value(HL_collection[0], -clip_value, clip_value) ) 19 | LH_collection[0] = tf.round( s * tf.clip_by_value(LH_collection[0], -clip_value, clip_value) ) 20 | HH_collection[0] = tf.round( s * tf.clip_by_value(HH_collection[0], -clip_value, clip_value) ) 21 | 22 | HL_collection[1] = tf.round( s * tf.clip_by_value(HL_collection[1], -clip_value, clip_value) ) 23 | LH_collection[1] = tf.round( s * tf.clip_by_value(LH_collection[1], -clip_value, clip_value) ) 24 | HH_collection[1] = tf.round( s * tf.clip_by_value(HH_collection[1], -clip_value, clip_value) ) 25 | 26 | HL_collection[2] = tf.round( s * tf.clip_by_value(HL_collection[2], -clip_value, clip_value) ) 27 | LH_collection[2] = tf.round( s * tf.clip_by_value(LH_collection[2], -clip_value, clip_value) ) 28 | HH_collection[2] = tf.round( s * tf.clip_by_value(HH_collection[2], -clip_value, clip_value) ) 29 | 30 | HL_collection[3] = tf.round( s * tf.clip_by_value(HL_collection[3], -clip_value, clip_value) ) 31 | LH_collection[3] = tf.round( s * tf.clip_by_value(LH_collection[3], -clip_value, clip_value) ) 32 | HH_collection[3] = tf.round( s * tf.clip_by_value(HH_collection[3], -clip_value, clip_value) ) 33 | 34 | LL = tf.round( s * tf.clip_by_value(LL, -clip_value, clip_value) ) 35 | 36 | 37 | return LL, HL_collection, LH_collection, HH_collection 38 | 39 | 40 | def de_quant(LL, HL_collection, LH_collection, HH_collection, s): 41 | 42 | HL_collection[0] = HL_collection[0] / s 43 | LH_collection[0] = LH_collection[0] / s 44 | HH_collection[0] = HH_collection[0] / s 45 | 46 | HL_collection[1] = HL_collection[1] / s 47 | LH_collection[1] = LH_collection[1] / s 48 | HH_collection[1] = HH_collection[1] / s 49 | 50 | HL_collection[2] = HL_collection[2] / s 51 | LH_collection[2] = LH_collection[2] / s 52 | HH_collection[2] = HH_collection[2] / s 53 | 54 | HL_collection[3] = HL_collection[3] / s 55 | LH_collection[3] = LH_collection[3] / s 56 | HH_collection[3] = HH_collection[3] / s 57 | 58 | LL = LL / s 59 | 60 | return LL, HL_collection, LH_collection, HH_collection -------------------------------------------------------------------------------- /code/lossy_multi_model/wavelet.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from __future__ import division 3 | from __future__ import print_function 4 | import argparse 5 | import sys 6 | import tensorflow as tf 7 | from math import log 8 | import numpy as np 9 | from PIL import Image 10 | import lifting97 11 | 12 | 13 | ############### one-level decomposition of lifting 97 14 | 15 | def decomposition(x): 16 | 17 | # step 1: for h 18 | L = x[:, 0::2, :, :] 19 | H = x[:, 1::2, :, :] 20 | 21 | L, H = lifting97.lifting97_forward(L, H) 22 | 23 | 24 | 25 | # step 2: for w, L 26 | 27 | L = tf.transpose(L, [0, 2, 1, 3]) 28 | 29 | LL = L[:, 0::2, :, :] 30 | HL = L[:, 1::2, :, :] 31 | 32 | LL, HL = lifting97.lifting97_forward(LL, HL) 33 | 34 | LL = tf.transpose(LL, [0, 2, 1, 3]) 35 | HL = tf.transpose(HL, [0, 2, 1, 3]) 36 | 37 | # step 2: for w, H 38 | 39 | H = tf.transpose(H, [0, 2, 1, 3]) 40 | 41 | LH = H[:, 0::2, :, :] 42 | HH = H[:, 1::2, :, :] 43 | 44 | LH, HH = lifting97.lifting97_forward(LH, HH) 45 | 46 | LH = tf.transpose(LH, [0, 2, 1, 3]) 47 | HH = tf.transpose(HH, [0, 2, 1, 3]) 48 | 49 | return LL, HL, LH, HH 50 | 51 | ############### one-level reconstruction of lifting 97 52 | 53 | def reconstruct_fun(up,bot): 54 | 55 | temp_L = tf.transpose(up, [0, 2, 1, 3]) 56 | temp_H = tf.transpose(bot, [0, 2, 1, 3]) 57 | 58 | x_shape = temp_L.get_shape().as_list() 59 | 60 | x_n = x_shape[0] 61 | x_h = x_shape[1] 62 | x_w = x_shape[2] 63 | x_c = x_shape[3] 64 | 65 | temp_L = tf.reshape(temp_L, [x_n, x_h * x_w, 1, x_c]) 66 | 67 | temp_H = tf.reshape(temp_H, [x_n, x_h * x_w, 1, x_c]) 68 | 69 | temp = tf.concat([temp_L,temp_H],2) 70 | 71 | temp = tf.reshape(temp,[x_n, x_h, 2*x_w, x_c]) 72 | 73 | recon = tf.transpose(temp, [0, 2, 1, 3]) 74 | 75 | return recon 76 | 77 | def reconstruct(LL,HL,LH,HH): 78 | 79 | LL = tf.transpose(LL, [0, 2, 1, 3]) 80 | HL = tf.transpose(HL, [0, 2, 1, 3]) 81 | 82 | LL, HL = lifting97.lifting97_inverse(LL, HL) 83 | 84 | L = reconstruct_fun(LL,HL) 85 | L = tf.transpose(L, [0, 2, 1, 3]) 86 | 87 | 88 | 89 | LH = tf.transpose(LH, [0, 2, 1, 3]) 90 | HH = tf.transpose(HH, [0, 2, 1, 3]) 91 | 92 | LH, HH = lifting97.lifting97_inverse(LH, HH) 93 | 94 | H = reconstruct_fun(LH, HH) 95 | H = tf.transpose(H, [0, 2, 1, 3]) 96 | 97 | 98 | 99 | L, H = lifting97.lifting97_inverse(L, H) 100 | 101 | recon = reconstruct_fun(L, H) 102 | 103 | 104 | return recon -------------------------------------------------------------------------------- /figs/Kodak.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mahaichuan/Versatile-Image-Compression/0973e7363f85c3eeb7865b270d14abcc37821708/figs/Kodak.PNG -------------------------------------------------------------------------------- /figs/Tecnick.PNG: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mahaichuan/Versatile-Image-Compression/0973e7363f85c3eeb7865b270d14abcc37821708/figs/Tecnick.PNG -------------------------------------------------------------------------------- /model/README.md: -------------------------------------------------------------------------------- 1 | [Lossy multi-model iWave++ is available now.](https://drive.google.com/drive/folders/10i0O66mOhE3C7AdG6-22kIFIIaRZ-nlj?usp=sharing) 2 | --------------------------------------------------------------------------------