test_video.py
官方固定版本源码 · 返回精读讲解 · 原始源码
1# Copyright (c) Microsoft Corporation.
2# Licensed under the MIT License.
3
4import argparse
5import math
6import os
7import concurrent.futures
8import multiprocessing
9import torch
10import json
11import numpy as np
12from PIL import Image
13from src.models.DCVC_net import DCVC_net
14from src.zoo.image import model_architectures as architectures
15import time
16from tqdm import tqdm
17import warnings
18from pytorch_msssim import ms_ssim
19
20
21warnings.filterwarnings("ignore", message="Setting attributes on ParameterList is not supported.")
22
23
24def str2bool(v):
25 return str(v).lower() in ("yes", "y", "true", "t", "1")
26
27
28def parse_args():
29 parser = argparse.ArgumentParser(description="Example testing script")
30
31 parser.add_argument('--i_frame_model_name', type=str, default="cheng2020-anchor")
32 parser.add_argument('--i_frame_model_path', type=str, nargs="+")
33 parser.add_argument('--model_path', type=str, nargs="+")
34 parser.add_argument('--test_config', type=str, required=True)
35 parser.add_argument("--worker", "-w", type=int, default=1, help="worker number")
36 parser.add_argument("--cuda", type=str2bool, nargs='?', const=True, default=False)
37 parser.add_argument("--cuda_device", default=None,
38 help="the cuda device used, e.g., 0; 0,1; 1,2,3; etc.")
39 parser.add_argument('--write_stream', type=str2bool, nargs='?',
40 const=True, default=False)
41 parser.add_argument("--write_recon_frame", type=str2bool,
42 nargs='?', const=True, default=False)
43 parser.add_argument('--recon_bin_path', type=str, default="recon_bin_path")
44 parser.add_argument('--output_json_result_path', type=str, required=True)
45 parser.add_argument("--model_type", type=str, default="psnr", help="psnr, msssim")
46
47 args = parser.parse_args()
48 return args
49
50
51def PSNR(input1, input2):
52 mse = torch.mean((input1 - input2) ** 2)
53 psnr = 20 * torch.log10(1 / torch.sqrt(mse))
54 return psnr.item()
55
56
57def read_frame_to_torch(path):
58 input_image = Image.open(path).convert('RGB')
59 input_image = np.asarray(input_image).astype('float64').transpose(2, 0, 1)
60 input_image = torch.from_numpy(input_image).type(torch.FloatTensor)
61 input_image = input_image.unsqueeze(0)/255
62 return input_image
63
64
65def write_torch_frame(frame, path):
66 frame_result = frame.clone()
67 frame_result = frame_result.cpu().detach().numpy().transpose(1, 2, 0)*255
68 frame_result = np.clip(np.rint(frame_result), 0, 255)
69 frame_result = Image.fromarray(frame_result.astype('uint8'), 'RGB')
70 frame_result.save(path)
71
72
73def encode_one(args_dict, device):
74 i_frame_load_checkpoint = torch.load(args_dict['i_frame_model_path'],
75 map_location=torch.device('cpu'))
76 i_frame_net = architectures[args_dict['i_frame_model_name']].from_state_dict(
77 i_frame_load_checkpoint).eval()
78
79 video_net = DCVC_net()
80 load_checkpoint = torch.load(args_dict['model_path'], map_location=torch.device('cpu'))
81 video_net.load_dict(load_checkpoint)
82
83 video_net = video_net.to(device)
84 video_net.eval()
85 i_frame_net = i_frame_net.to(device)
86 i_frame_net.eval()
87 if args_dict['write_stream']:
88 video_net.update(force=True)
89 i_frame_net.update(force=True)
90
91 sub_dir_name = args_dict['video_path']
92 ref_frame = None
93 frame_types = []
94 qualitys = []
95 bits = []
96 bits_mv_y = []
97 bits_mv_z = []
98 bits_y = []
99 bits_z = []
100
101 gop_size = args_dict['gop']
102 frame_pixel_num = 0
103 frame_num = args_dict['frame_num']
104
105 recon_bin_folder = os.path.join(args_dict['recon_bin_path'], sub_dir_name,
106 os.path.basename(args_dict['model_path'])[:-4])
107 if not os.path.exists(recon_bin_folder):
108 os.makedirs(recon_bin_folder)
109
110 # Figure out the naming convention
111 pngs = os.listdir(os.path.join(args_dict['dataset_path'], sub_dir_name))
112 if 'im1.png' in pngs:
113 padding = 1
114 elif 'im00001.png' in pngs:
115 padding = 5
116 else:
117 raise ValueError('unknown image naming convention; please specify')
118
119 with torch.no_grad():
120 for frame_idx in range(frame_num):
121 ori_frame = read_frame_to_torch(
122 os.path.join(args_dict['dataset_path'],
123 sub_dir_name,
124 f"im{str(frame_idx+1).zfill(padding)}.png"))
125 ori_frame = ori_frame.to(device)
126
127 if frame_pixel_num == 0:
128 frame_pixel_num = ori_frame.shape[2]*ori_frame.shape[3]
129 else:
130 assert(frame_pixel_num == ori_frame.shape[2]*ori_frame.shape[3])
131
132 if args_dict['write_stream']:
133 bin_path = os.path.join(recon_bin_folder, f"{frame_idx}.bin")
134 if frame_idx % gop_size == 0:
135 result = i_frame_net.encode_decode(ori_frame, bin_path)
136 ref_frame = result["x_hat"]
137 bpp = result["bpp"]
138 frame_types.append(0)
139 bits.append(bpp*frame_pixel_num)
140 bits_mv_y.append(0)
141 bits_mv_z.append(0)
142 bits_y.append(0)
143 bits_z.append(0)
144 else:
145 result = video_net.encode_decode(ref_frame, ori_frame, bin_path)
146 ref_frame = result['recon_image']
147 bpp = result['bpp']
148 frame_types.append(1)
149 bits.append(bpp*frame_pixel_num)
150 bits_mv_y.append(result['bpp_mv_y']*frame_pixel_num)
151 bits_mv_z.append(result['bpp_mv_z']*frame_pixel_num)
152 bits_y.append(result['bpp_y']*frame_pixel_num)
153 bits_z.append(result['bpp_z']*frame_pixel_num)
154 else:
155 if frame_idx % gop_size == 0:
156 result = i_frame_net(ori_frame)
157 bit = sum((torch.log(likelihoods).sum() / (-math.log(2)))
158 for likelihoods in result["likelihoods"].values())
159 ref_frame = result["x_hat"]
160 frame_types.append(0)
161 bits.append(bit.item())
162 bits_mv_y.append(0)
163 bits_mv_z.append(0)
164 bits_y.append(0)
165 bits_z.append(0)
166 else:
167 result = video_net(ref_frame, ori_frame)
168 ref_frame = result['recon_image']
169 bpp = result['bpp']
170 frame_types.append(1)
171 bits.append(bpp.item()*frame_pixel_num)
172 bits_mv_y.append(result['bpp_mv_y'].item()*frame_pixel_num)
173 bits_mv_z.append(result['bpp_mv_z'].item()*frame_pixel_num)
174 bits_y.append(result['bpp_y'].item()*frame_pixel_num)
175 bits_z.append(result['bpp_z'].item()*frame_pixel_num)
176
177 ref_frame = ref_frame.clamp_(0, 1)
178 if args_dict['write_recon_frame']:
179 write_torch_frame(ref_frame.squeeze(), os.path.join(recon_bin_folder, f"recon_frame_{frame_idx}.png"))
180 if args_dict['model_type'] == 'psnr':
181 qualitys.append(PSNR(ref_frame, ori_frame))
182 else:
183 qualitys.append(
184 ms_ssim(ref_frame, ori_frame, data_range=1.0).item())
185
186 cur_all_i_frame_bit = 0
187 cur_all_i_frame_quality = 0
188 cur_all_p_frame_bit = 0
189 cur_all_p_frame_bit_mv_y = 0
190 cur_all_p_frame_bit_mv_z = 0
191 cur_all_p_frame_bit_y = 0
192 cur_all_p_frame_bit_z = 0
193 cur_all_p_frame_quality = 0
194 cur_i_frame_num = 0
195 cur_p_frame_num = 0
196 for idx in range(frame_num):
197 if frame_types[idx] == 0:
198 cur_all_i_frame_bit += bits[idx]
199 cur_all_i_frame_quality += qualitys[idx]
200 cur_i_frame_num += 1
201 else:
202 cur_all_p_frame_bit += bits[idx]
203 cur_all_p_frame_bit_mv_y += bits_mv_y[idx]
204 cur_all_p_frame_bit_mv_z += bits_mv_z[idx]
205 cur_all_p_frame_bit_y += bits_y[idx]
206 cur_all_p_frame_bit_z += bits_z[idx]
207 cur_all_p_frame_quality += qualitys[idx]
208 cur_p_frame_num += 1
209
210 log_result = {}
211 log_result['name'] = f"{os.path.basename(args_dict['model_path'])}_{sub_dir_name}"
212 log_result['ds_name'] = args_dict['ds_name']
213 log_result['video_path'] = args_dict['video_path']
214 log_result['frame_pixel_num'] = frame_pixel_num
215 log_result['i_frame_num'] = cur_i_frame_num
216 log_result['p_frame_num'] = cur_p_frame_num
217 log_result['ave_i_frame_bpp'] = cur_all_i_frame_bit / cur_i_frame_num / frame_pixel_num
218 log_result['ave_i_frame_quality'] = cur_all_i_frame_quality / cur_i_frame_num
219 if cur_p_frame_num > 0:
220 total_p_pixel_num = cur_p_frame_num * frame_pixel_num
221 log_result['ave_p_frame_bpp'] = cur_all_p_frame_bit / total_p_pixel_num
222 log_result['ave_p_frame_bpp_mv_y'] = cur_all_p_frame_bit_mv_y / total_p_pixel_num
223 log_result['ave_p_frame_bpp_mv_z'] = cur_all_p_frame_bit_mv_z / total_p_pixel_num
224 log_result['ave_p_frame_bpp_y'] = cur_all_p_frame_bit_y / total_p_pixel_num
225 log_result['ave_p_frame_bpp_z'] = cur_all_p_frame_bit_z / total_p_pixel_num
226 log_result['ave_p_frame_quality'] = cur_all_p_frame_quality / cur_p_frame_num
227 else:
228 log_result['ave_p_frame_bpp'] = 0
229 log_result['ave_p_frame_quality'] = 0
230 log_result['ave_p_frame_bpp_mv_y'] = 0
231 log_result['ave_p_frame_bpp_mv_z'] = 0
232 log_result['ave_p_frame_bpp_y'] = 0
233 log_result['ave_p_frame_bpp_z'] = 0
234 log_result['ave_all_frame_bpp'] = (cur_all_i_frame_bit + cur_all_p_frame_bit) / \
235 (frame_num * frame_pixel_num)
236 log_result['ave_all_frame_quality'] = (cur_all_i_frame_quality + cur_all_p_frame_quality) / frame_num
237 return log_result
238
239
240def worker(use_cuda, args):
241 if args['write_stream']:
242 torch.backends.cudnn.benchmark = False
243 if 'use_deterministic_algorithms' in dir(torch):
244 torch.use_deterministic_algorithms(True)
245 else:
246 torch.set_deterministic(True)
247 torch.manual_seed(0)
248 torch.set_num_threads(1)
249 np.random.seed(seed=0)
250 gpu_num = 0
251 if use_cuda:
252 gpu_num = torch.cuda.device_count()
253
254 process_name = multiprocessing.current_process().name
255 process_idx = int(process_name[process_name.rfind('-') + 1:])
256 gpu_id = -1
257 if gpu_num > 0:
258 gpu_id = process_idx % gpu_num
259 if gpu_id >= 0:
260 device = f"cuda:{gpu_id}"
261 else:
262 device = "cpu"
263
264 result = encode_one(args, device)
265 result['model_idx'] = args['model_idx']
266 return result
267
268
269def filter_dict(result):
270 keys = ['i_frame_num', 'p_frame_num', 'ave_i_frame_bpp', 'ave_i_frame_quality', 'ave_p_frame_bpp',
271 'ave_p_frame_bpp_mv_y', 'ave_p_frame_bpp_mv_z', 'ave_p_frame_bpp_y',
272 'ave_p_frame_bpp_z', 'ave_p_frame_quality', 'ave_all_frame_bpp', 'ave_all_frame_quality']
273 res = {k: v for k, v in result.items() if k in keys}
274 return res
275
276
277def main():
278 torch.backends.cudnn.enabled = True
279 args = parse_args()
280 if args.cuda_device is not None and args.cuda_device != '':
281 os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda_device
282 os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"
283 worker_num = args.worker
284 assert worker_num >= 1
285
286 with open(args.test_config) as f:
287 config = json.load(f)
288
289 multiprocessing.set_start_method("spawn")
290 threadpool_executor = concurrent.futures.ProcessPoolExecutor(max_workers=worker_num)
291 objs = []
292
293 count_frames = 0
294 count_sequences = 0
295 begin_time = time.time()
296 for ds_name in config:
297 for seq_name in config[ds_name]['sequences']:
298 count_sequences += 1
299 for model_idx in range(len(args.model_path)):
300 cur_dict = {}
301 cur_dict['model_idx'] = model_idx
302 cur_dict['i_frame_model_path'] = args.i_frame_model_path[model_idx]
303 cur_dict['i_frame_model_name'] = args.i_frame_model_name
304 cur_dict['model_path'] = args.model_path[model_idx]
305 cur_dict['video_path'] = seq_name
306 cur_dict['gop'] = config[ds_name]['sequences'][seq_name]['gop']
307 cur_dict['frame_num'] = config[ds_name]['sequences'][seq_name]['frames']
308 cur_dict['dataset_path'] = config[ds_name]['base_path']
309 cur_dict['write_stream'] = args.write_stream
310 cur_dict['write_recon_frame'] = args.write_recon_frame
311 cur_dict['recon_bin_path'] = args.recon_bin_path
312 cur_dict['model_type'] = args.model_type
313 cur_dict['ds_name'] = ds_name
314
315 count_frames += cur_dict['frame_num']
316
317 obj = threadpool_executor.submit(
318 worker,
319 args.cuda,
320 cur_dict)
321 objs.append(obj)
322
323 results = []
324 for obj in tqdm(objs):
325 result = obj.result()
326 results.append(result)
327
328 log_result = {}
329
330 for ds_name in config:
331 log_result[ds_name] = {}
332 for seq in config[ds_name]['sequences']:
333 log_result[ds_name][seq] = {}
334 for model_idx in range(len(args.model_path)):
335 ckpt = os.path.basename(args.model_path[model_idx])
336 for res in results:
337 if res['name'].startswith(ckpt) and ds_name == res['ds_name'] \
338 and seq == res['video_path']:
339 log_result[ds_name][seq][ckpt] = filter_dict(res)
340
341 with open(args.output_json_result_path, 'w') as fp:
342 json.dump(log_result, fp, indent=2)
343
344 total_minutes = (time.time() - begin_time) / 60
345
346 count_models = len(args.model_path)
347 count_frames = count_frames // count_models
348 print('Test finished')
349 print(f'Tested {count_models} models on {count_frames} frames from {count_sequences} sequences')
350 print(f'Total elapsed time: {total_minutes:.1f} min')
351
352
353if __name__ == "__main__":
354 main()