add rep resuklts on kodak dataset

Rep_results/cal_results.py

@@ -0,0 +1,51 @@
+import argparse
+import glob
+import os
+import csv
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--exp_dir', help='exp_dir_of the results of each png file', default='Rep_results/kodak/10w_5loop_train')
+args = parser.parse_args()
+
+def cal_avg_psnr_bpp(inputpath):
+
+    f = open(inputpath,"r",encoding='utf-8')
+    lines =f.readlines()
+    datas = lines[1].split()
+    f.close()
+    return float(datas[3]),float(datas[6])
+
+def takeOne(elem):
+    return elem[0]
+
+
+if __name__ == '__main__':
+    lambda_dirs = glob.glob(os.path.join(args.exp_dir,'lamda*'),recursive=False)
+    for lambda_dir in lambda_dirs:
+        encoded_images = glob.glob(os.path.join(lambda_dir,'*/results_best.tsv'))
+        results = {'avg_bpp': [], 'avg_psnr': []}
+        avg_psnr = avg_bpp = 0.0
+        img_nums = len(encoded_images)
+        print('Images count: ' + str(img_nums))
+
+
+        writer_list = [] # 'name' 'psnr' 'bpp'
+        for i,img in enumerate(encoded_images):
+            psnr,bpp = cal_avg_psnr_bpp(img)
+            name = img.split('/')[-2]
+            writer_list.append([name,psnr,bpp])
+            avg_psnr += psnr
+            avg_bpp += bpp
+        avg_psnr /= img_nums
+        avg_bpp /= img_nums
+
+        writer_list.sort(key=takeOne)
+
+
+        with open(os.path.join(args.exp_dir,lambda_dir.split('/')[-1] +'_encoder_results.csv'), 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['name','psnr','bpp'])
+            writer.writerows(writer_list)
+            writer.writerow(['','avg_psnr','avg_bpp'])
+            writer.writerow(['',avg_psnr, avg_bpp])
+        f.close()

Rep_results/encode_one_dataset.py

@@ -0,0 +1,86 @@
+import glob
+import os
+import subprocess
+import argparse
+
+POSSIBLE_DATASETS = ['kodak', 'clic20-pro-valid', 'clic22-test', 'jvet', 'vcip2023_4k','vcip_4k_g1','vcip_4k_g2','vcip_4k_g3','vcip_4k_g4']
+bpp_offical = {
+    'I01':0.0015,
+    'I02':0.0009,
+    'I03':0.0009,
+    'I04':0.0005,
+    'I05':0.0005,
+    'I06':0.0009,
+    'I07':0.0015,
+    'I08':0.0003,
+    'I09':0.0007,
+    'I10':0.0009,
+    'I11':0.0009,
+    'I12':0.00085,
+    'I13':0.0002,
+    'I14':0.001,
+    'I15':0.00085,
+    'I16':0.00075,
+    'I17':0.0007,
+    'I18':0.0008,
+    'I19':0.0007,
+    'I20':0.0007,
+}
+
+parser = argparse.ArgumentParser()
+parser.add_argument('dataset_name', help=f'Possible values {POSSIBLE_DATASETS}.')
+parser.add_argument('--lamda', help='d+lamda*R',default=0.0001)
+parser.add_argument('--lamda_list', type=list, default=[0.00012,0.0008,0.0008,0.0008,0.00012,0.0008])
+
+args = parser.parse_args()
+
+def lamda_to_str(lamda) :
+    str_lamda = list()
+    tmp = lamda
+    while tmp< 1 :
+        str_lamda.append(int(tmp * 10))
+        tmp = tmp * 10
+    return str_lamda
+
+assert args.dataset_name in POSSIBLE_DATASETS, \
+    f'Argument must be in {POSSIBLE_DATASETS}. Found {args.dataset_name}!'
+
+
+# encode images one by one
+current_dir_path = os.path.dirname(__file__)
+dataset_path = '/home/hangyu.li/workspace/Cool-Chic/dataset'
+encoded_image_path = os.path.join(dataset_path,args.dataset_name)
+
+cool_chic_encode_path = os.path.join(current_dir_path, '../src/encode.py')
+
+
+encoded_images = glob.glob(os.path.join(encoded_image_path,'*.png'))
+for encoded_img in encoded_images :
+    imgname = encoded_img.split("/")[-1].split(".")[0]
+    output_dir = os.path.join(current_dir_path, args.dataset_name, '10w_5loop_train',f"lamda-{args.lamda}", imgname)
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+    output_bitstream_path = os.path.join(output_dir, f'{imgname}.bin')
+    model_save_path = os.path.join(output_dir, f'video_encoder.pt')
+    enc_results_path = os.path.join(output_dir, f'{imgname}_encoder_results.txt')
+    qstep_results_pkl = os.path.join(output_dir, f'{imgname}_qstep_results.pkl')
+    if os.path.exists(model_save_path):
+        print(f'skipped_{imgname}')
+        continue
+
+
+    print(f'\nencoding image: {encoded_img}')
+    cmd = f'CUDA_VISIBLE_DEVICES=2 python3 {cool_chic_encode_path} \
+        --input {encoded_img} \
+        --output {output_bitstream_path} \
+        --workdir={output_dir}/ \
+        --lmbda={args.lamda}\
+        --start_lr=1e-2\
+        --layers_synthesis=40-1-linear-relu,3-1-linear-relu,X-3-residual-relu,X-3-residual-none\
+        --upsampling_kernel_size=8                      \
+        --layers_arm=24,24\
+        --n_ctx_rowcol=3\
+        --n_ft_per_res=1,1,1,1,1,1,1                    \
+        --n_itr=100000                                    \
+        --n_train_loops=5'
+    subprocess.call(cmd, shell=True)

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/archi.txt

@@ -0,0 +1,90 @@
+CoolChicEncoder(
+  (latent_grids): ParameterList(
+      (0): Parameter containing: [torch.float32 of size 1x1x512x768]
+      (1): Parameter containing: [torch.float32 of size 1x1x256x384]
+      (2): Parameter containing: [torch.float32 of size 1x1x128x192]
+      (3): Parameter containing: [torch.float32 of size 1x1x64x96]
+      (4): Parameter containing: [torch.float32 of size 1x1x32x48]
+      (5): Parameter containing: [torch.float32 of size 1x1x16x24]
+      (6): Parameter containing: [torch.float32 of size 1x1x8x12]
+  )
+  (synthesis): Synthesis(
+    (layers): Sequential(
+      (0): SynthesisLayer(
+        (pad): ReplicationPad2d((0, 0, 0, 0))
+        (conv_layer): Conv2d(7, 40, kernel_size=(1, 1), stride=(1, 1))
+        (non_linearity): ReLU()
+      )
+      (1): SynthesisLayer(
+        (pad): ReplicationPad2d((0, 0, 0, 0))
+        (conv_layer): Conv2d(40, 3, kernel_size=(1, 1), stride=(1, 1))
+        (non_linearity): ReLU()
+      )
+      (2): SynthesisResidualLayer(
+        (pad): ReplicationPad2d((1, 1, 1, 1))
+        (conv_layer): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
+        (non_linearity): ReLU()
+      )
+      (3): SynthesisResidualLayer(
+        (pad): ReplicationPad2d((1, 1, 1, 1))
+        (conv_layer): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
+        (non_linearity): Identity()
+      )
+    )
+  )
+  (noise_quantizer): NoiseQuantizer()
+  (ste_quantizer): STEQuantizer()
+  (upsampling): Upsampling(
+    (upsampling_layer): ConvTranspose2d(1, 1, kernel_size=(1, 1, 8, 8), stride=(2, 2), bias=False)
+  )
+  (arm): Arm(
+    (mlp): Sequential(
+      (0): CustomLinearResBlock()
+      (1): ReLU()
+      (2): CustomLinearResBlock()
+      (3): ReLU()
+      (4): CustomLinear()
+    )
+  )
+)
+
+| module                                  | #parameters or shape   | #flops    |
+|:----------------------------------------|:-----------------------|:----------|
+| model                                   | 0.526M                 | 0.901G    |
+|  latent_grids                           |  0.524M                |           |
+|   latent_grids.0                        |   (1, 1, 512, 768)     |           |
+|   latent_grids.1                        |   (1, 1, 256, 384)     |           |
+|   latent_grids.2                        |   (1, 1, 128, 192)     |           |
+|   latent_grids.3                        |   (1, 1, 64, 96)       |           |
+|   latent_grids.4                        |   (1, 1, 32, 48)       |           |
+|   latent_grids.5                        |   (1, 1, 16, 24)       |           |
+|   latent_grids.6                        |   (1, 1, 8, 12)        |           |
+|  synthesis.layers                       |  0.611K                |  0.221G   |
+|   synthesis.layers.0.conv_layer         |   0.32K                |   0.11G   |
+|    synthesis.layers.0.conv_layer.weight |    (40, 7, 1, 1)       |           |
+|    synthesis.layers.0.conv_layer.bias   |    (40,)               |           |
+|   synthesis.layers.1.conv_layer         |   0.123K               |   47.186M |
+|    synthesis.layers.1.conv_layer.weight |    (3, 40, 1, 1)       |           |
+|    synthesis.layers.1.conv_layer.bias   |    (3,)                |           |
+|   synthesis.layers.2.conv_layer         |   84                   |   31.85M  |
+|    synthesis.layers.2.conv_layer.weight |    (3, 3, 3, 3)        |           |
+|    synthesis.layers.2.conv_layer.bias   |    (3,)                |           |
+|   synthesis.layers.3.conv_layer         |   84                   |   31.85M  |
+|    synthesis.layers.3.conv_layer.weight |    (3, 3, 3, 3)        |           |
+|    synthesis.layers.3.conv_layer.bias   |    (3,)                |           |
+|  upsampling.upsampling_layer            |  64                    |  50.909M  |
+|   upsampling.upsampling_layer.weight    |   (1, 1, 8, 8)         |           |
+|  arm.mlp                                |  1.25K                 |  0.629G   |
+|   arm.mlp.0                             |   0.6K                 |   0.302G  |
+|    arm.mlp.0.weight                     |    (24, 24)            |           |
+|    arm.mlp.0.bias                       |    (24,)               |           |
+|   arm.mlp.2                             |   0.6K                 |   0.302G  |
+|    arm.mlp.2.weight                     |    (24, 24)            |           |
+|    arm.mlp.2.bias                       |    (24,)               |           |
+|   arm.mlp.4                             |   50                   |   25.164M |
+|    arm.mlp.4.weight                     |    (2, 24)             |           |
+|    arm.mlp.4.bias                       |    (2,)                |           |
+
+----------------------------------
+Total MAC / decoded pixel: 2291.4
+----------------------------------

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_best.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.298877   0.040827   2.119735   40.818605  25.542143  -0.000000  2.160562   0.000100   4042.056376 104884     2291.371094 1.965579   0.069934   0.042715   0.029734   0.007691   0.002903   0.001179   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.021354   0.001385   0.018087   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_loop_1.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.298877   0.040827   2.119735   40.818605  25.542143  -0.000000  2.160562   0.000100   4042.056376 104884     2291.371094 1.965579   0.069934   0.042715   0.029734   0.007691   0.002903   0.001179   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.021354   0.001385   0.018087   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_loop_2.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.316812   0.037137   2.126358   39.979968  24.795748  -0.000000  2.163494   0.000100   8007.482905 210248     2291.371094 1.916443   0.134619   0.043232   0.019158   0.007753   0.003865   0.001290   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.020064   0.001112   0.015961   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_loop_3.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.319148   0.036207   2.129885   39.891109  24.616515  -0.000000  2.166092   0.000100   11986.141246 315583     2291.371094 1.930957   0.112920   0.047470   0.024197   0.010584   0.002664   0.001093   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.020020   0.001078   0.015109   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_loop_4.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.300215   0.039554   2.123958   40.764237  25.547711  -0.000000  2.163512   0.000100   15966.292464 420585     2291.371094 1.975977   0.076771   0.032639   0.026006   0.007967   0.003330   0.001269   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.020667   0.001397   0.017490   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/frame_000/results_loop_5.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.314463   0.037093   2.139757   40.142254  24.893925  -0.000000  2.176850   0.000100   19952.731834 525674     2291.371094 1.934386   0.111429   0.042555   0.035368   0.011020   0.003671   0.001329   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.020375   0.001572   0.015146   

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/param.txt

@@ -0,0 +1 @@
+['/home/hangyu.li/workspace/Cool-Chic/Rep_results/../src/encode.py', '--input', '/home/hangyu.li/workspace/Cool-Chic/dataset/kodak/kodim01.png', '--output', '/home/hangyu.li/workspace/Cool-Chic/Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/kodim01.bin', '--workdir=/home/hangyu.li/workspace/Cool-Chic/Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/', '--lmbda=0.0001', '--start_lr=1e-2', '--layers_synthesis=40-1-linear-relu,3-1-linear-relu,X-3-residual-relu,X-3-residual-none', '--upsampling_kernel_size=8', '--layers_arm=24,24', '--n_ctx_rowcol=3', '--n_ft_per_res=1,1,1,1,1,1,1', '--n_itr=100000', '--n_train_loops=5']

Rep_results/kodak/10w_5loop_train/lamda-0.0001/kodim01/results_best.tsv

@@ -0,0 +1,2 @@
+loss       nn_bpp     latent_bpp psnr_db    ms_ssim_db lpips_db   total_rate_bpp lmbda      time_sec   itr        mac_decoded_pixel feature_rate_bpp_00 feature_rate_bpp_01 feature_rate_bpp_02 feature_rate_bpp_03 feature_rate_bpp_04 feature_rate_bpp_05 feature_rate_bpp_06 alpha      beta       pred_db    dummy_pred img_size   n_pixels   display_order coding_order seq_name   arm_rate_bpp upsampling_rate_bpp synthesis_rate_bpp 
+0.298877   0.040827   2.119735   40.818605  25.542143  -0.000000  2.160562   0.000100   4042.056376 104884     2291.371094 1.965579   0.069934   0.042715   0.029734   0.007691   0.002903   0.001179   1.000000   1.000000   7.015595   nan        512x768    393216     0          0          kodim01    0.021354   0.001385   0.018087