test_video.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import argparse
import concurrent.futures
import io
import json
import multiprocessing
import numpy as np
import os
import time
import torch

from tqdm import tqdm

from src.models.image_model import DMCI
from src.utils.common import ModelStructure, str2bool, create_folder, generate_log_json, \
    get_state_dict, dump_json, set_torch_env
from src.utils.metrics import calc_msssim, calc_msssim_rgb, calc_psnr
from src.utils.stream_helper import SPSHelper, NalType, write_sps, read_header, \
    read_sps_remaining, read_ip_remaining, write_ip
from src.utils.transforms import rgb2ycbcr, ycbcr2rgb, yuv_444_to_420, ycbcr420_to_444_np
from src.utils.video_reader import PNGReader, YUV420Reader
from src.utils.video_writer import PNGWriter, YUV420Writer


def finalize_model(net, device):
    net = net.half().to(device)
    return net.to(memory_format=torch.channels_last)


def get_distortion(args, x_hat, y, u, v, rgb):
    x_hat = x_hat + 0.5
    if args['src_type'] == 'yuv420':
        y_rec, uv_rec = yuv_444_to_420(x_hat)
        y_rec = torch.clamp(y_rec * 255, 0, 255).squeeze(0).cpu().numpy()
        uv_rec = torch.clamp(uv_rec * 255, 0, 255).squeeze(0).cpu().numpy()
        y_rec = y_rec[0, :, :]
        u_rec = uv_rec[0, :, :]
        v_rec = uv_rec[1, :, :]
        psnr_y = calc_psnr(y, y_rec)
        psnr_u = calc_psnr(u, u_rec)
        psnr_v = calc_psnr(v, v_rec)
        psnr = (6 * psnr_y + psnr_u + psnr_v) / 8
        if args['calc_ssim']:
            ssim_y = calc_msssim(y, y_rec)
            ssim_u = calc_msssim(u, u_rec)
            ssim_v = calc_msssim(v, v_rec)
        else:
            ssim_y, ssim_u, ssim_v = 0., 0., 0.
        ssim = (6 * ssim_y + ssim_u + ssim_v) / 8

        curr_psnr = [psnr, psnr_y, psnr_u, psnr_v]
        curr_ssim = [ssim, ssim_y, ssim_u, ssim_v]
    else:
        assert args['src_type'] == 'png'
        rgb_rec = ycbcr2rgb(x_hat)
        rgb_rec = torch.clamp(rgb_rec * 255, 0, 255).squeeze(0).cpu().numpy()
        psnr = calc_psnr(rgb, rgb_rec)
        if args['calc_ssim']:
            msssim = calc_msssim_rgb(rgb, rgb_rec)
        else:
            msssim = 0.
        curr_psnr = [psnr]
        curr_ssim = [msssim]
    return curr_psnr, curr_ssim


def get_src_frame(args, src_reader, device, maximum_read, is_intra, np_only=False):
    if args['src_type'] == 'yuv420':
        processed = 0
        x, y, u, v, rgb = [], [], [], [], []
        while processed < maximum_read:
            curr_y, curr_uv = src_reader.read_one_frame()
            curr_yuv = ycbcr420_to_444_np(curr_y, curr_uv)
            curr_x = torch.from_numpy(curr_yuv).unsqueeze(0)
            curr_y = curr_y[0, :, :]
            curr_u = curr_uv[0, :, :]
            curr_v = curr_uv[1, :, :]
            curr_rgb = None
            x.append(curr_x)
            y.append(curr_y)
            u.append(curr_u)
            v.append(curr_v)
            rgb.append(curr_rgb)
            processed += 1
    else:
        assert args['src_type'] == 'png'
        processed = 0
        x, y, u, v, rgb = [], [], [], [], []
        while processed < maximum_read:
            curr_rgb = src_reader.read_one_frame()
            curr_x = torch.from_numpy(curr_rgb).unsqueeze(0).to(device=device, non_blocking=True)
            curr_x = curr_x.float() / 255.0
            curr_x = rgb2ycbcr(curr_x)
            curr_y, curr_u, curr_v = None, None, None
            x.append(curr_x)
            y.append(curr_y)
            u.append(curr_u)
            v.append(curr_v)
            rgb.append(curr_rgb)
            processed += 1

    while not is_intra and processed < g_frame_delay:
        x.append(x[-1])
        y.append(y[-1])
        u.append(u[-1])
        v.append(v[-1])
        rgb.append(rgb[-1])
        processed += 1

    if np_only:
        return y, u, v, rgb

    x = torch.cat(x, dim=1)
    if args['src_type'] == 'yuv420':
        x = x.to(device=device, non_blocking=True)
    x = x.half()
    if args['src_type'] == 'yuv420':
        x = x / 255.0
    x = x - 0.5
    x = x.to(memory_format=torch.channels_last)
    return x, y, u, v, rgb


def get_src_reader(args):
    if args['src_type'] == 'png':
        src_reader = PNGReader(args['src_path'], args['src_width'], args['src_height'])
    elif args['src_type'] == 'yuv420':
        src_reader = YUV420Reader(args['src_path'], args['src_width'], args['src_height'])
    else:
        assert False
    return src_reader


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument('--skip_thres', type=float, default=0)
    parser.add_argument('--model_path_i', type=str)
    parser.add_argument('--model_path_p',  type=str)
    parser.add_argument('--rate_num', type=int, default=4)
    parser.add_argument('--qp_i', type=int, nargs='+')
    parser.add_argument('--qp_p', type=int, nargs='+')
    parser.add_argument('--force_intra', type=str2bool, default=False)
    parser.add_argument('--force_frame_num', type=int, default=-1)
    parser.add_argument('--force_intra_period', type=int, default=-1)
    parser.add_argument('--reset_interval', type=int, default=32)
    parser.add_argument('--test_config', type=str, required=True)
    parser.add_argument('--force_root_path', type=str, default=None, required=False)
    parser.add_argument('--worker', '-w', type=int, default=1, help='worker number')
    parser.add_argument('--model_structure', type=str, default='ld', choices=['htl', 'hts', 'ld'])
    parser.add_argument('--cuda_idx', type=int, nargs='+', help='GPU indexes to use')
    parser.add_argument('--calc_ssim', type=str2bool, default=False, required=False)
    parser.add_argument('--check_existing', type=str2bool, default=False)
    parser.add_argument('--stream_path', type=str, default='out_bin')
    parser.add_argument('--save_decoded_frame', type=str2bool, default=False)
    parser.add_argument('--output_path', type=str, required=True)
    parser.add_argument('--verbose_json', type=str2bool, default=False)
    parser.add_argument('--verbose', type=int, default=0)

    args = parser.parse_args()
    return args


@torch.inference_mode()
def run_one_point_with_stream(p_net, i_net, args):
    if args['check_existing'] and os.path.exists(args['curr_json_path']) and \
            os.path.exists(args['curr_bin_path']):
        with open(args['curr_json_path']) as f:
            log_result = json.load(f)
            if log_result['i_frame_num'] + log_result['p_frame_num'] == args['frame_num']:
                return log_result
            print(f'incorrect log for {args['curr_json_path']}, try to rerun.')

    frame_num = args['frame_num']
    save_decoded_frame = args['save_decoded_frame']
    verbose = args['verbose']
    reset_interval = args['reset_interval']
    intra_period = args['intra_period']
    verbose_json = args['verbose_json']
    device = next(i_net.parameters()).device

    src_reader = get_src_reader(args)
    pic_height = args['src_height']
    pic_width = args['src_width']
    padding_r, padding_b = DMCI.get_padding_size(pic_height, pic_width, 16)

    frame_types = []
    psnrs = []
    msssims = []
    bits = []

    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    start_time = time.time()
    encoding_time = []
    decoding_time = []

    output_buff = io.BytesIO()
    sps_helper = SPSHelper()

    frame_idx = 0
    while frame_idx < frame_num:
        is_intra = False
        if frame_idx == 0 or intra_period == 1:
            is_intra = True
        if intra_period > 1 and frame_idx != 1:
            assert intra_period % g_frame_delay == 0
            if frame_idx % intra_period == 1:
                is_intra = True

        maximum_read = min(g_frame_delay, frame_num - frame_idx)
        if is_intra:
            maximum_read = 1

        x, y, u, v, rgb = get_src_frame(args, src_reader, device, maximum_read, is_intra)

        torch.cuda.synchronize(device=device)
        start_event.record()

        if is_intra:
            curr_qp = args['qp_i']
            reset_feature_memory = 0
            encoded = i_net.compress(x, curr_qp, padding_b, padding_r)
            if not args['force_intra']:
                p_net.clear_dpb()
                p_net.add_ref_feature_from_frame(encoded['x_hat'])
            frame_types.append(0)
        else:
            if reset_interval > 0 and (frame_idx + g_frame_delay) % reset_interval == 1:
                reset_feature_memory = 1
            else:
                reset_feature_memory = 0
            curr_qp = args['qp_p']

            encoded = p_net.compress(x, curr_qp, reset_feature_memory, padding_b, padding_r)
            for _ in range(maximum_read):
                frame_types.append(1)

        sps = {
            'sps_id': -1,
            'height': pic_height,
            'width': pic_width,
        }
        sps_id, sps_new = sps_helper.get_sps_id(sps)
        sps['sps_id'] = sps_id
        sps_bytes = 0
        if sps_new:
            sps_bytes = write_sps(output_buff, sps)
            if verbose >= 2:
                print('new sps', sps)
        stream_bytes = write_ip(output_buff, is_intra, sps_id, curr_qp,
                                encoded['ec_parallel'], reset_feature_memory,
                                encoded['bit_stream'])
        bits.append(stream_bytes * 8 + sps_bytes * 8)
        for _ in range(1, maximum_read):
            bits.append(0)

        end_event.record()
        torch.cuda.synchronize(device=device)

        frame_time = start_event.elapsed_time(end_event) / 1000.0
        encoding_time.append(frame_time)

        if verbose >= 2:
            print(f'frame {frame_idx} encoded, {frame_time * 1000:.3f} ms, '
                  f'bits: {bits[-maximum_read]}')

        frame_idx += maximum_read

    src_reader.close()
    with open(args['curr_bin_path'], 'wb') as output_file:
        bytes_buffer = output_buff.getbuffer()
        output_file.write(bytes_buffer)
        total_bytes = bytes_buffer.nbytes
        bytes_buffer.release()
    total_kbps = int(total_bytes * 8 / (frame_num / 30) / 1000)  # assume 30 fps
    output_buff.close()
    sps_helper = SPSHelper()
    with open(args['curr_bin_path'], 'rb') as input_file:
        input_buff = io.BytesIO(input_file.read())
    decoded_frame_number = 0
    src_reader = get_src_reader(args)

    if save_decoded_frame:
        if args['src_type'] == 'png':
            recon_writer = PNGWriter(args['bin_folder'], args['src_width'], args['src_height'])
        elif args['src_type'] == 'yuv420':
            output_yuv_path = args['curr_rec_path'].replace('.yuv', f'_{total_kbps}kbps.yuv')
            recon_writer = YUV420Writer(output_yuv_path, args['src_width'], args['src_height'])

    while decoded_frame_number < frame_num:
        torch.cuda.synchronize(device=device)
        start_event.record()

        header = read_header(input_buff)
        while header['nal_type'] == NalType.NAL_SPS:
            sps = read_sps_remaining(input_buff, header['sps_id'])
            sps_helper.add_sps_by_id(sps)
            if verbose >= 2:
                print('new sps', sps)
            header = read_header(input_buff)
            continue
        sps_id = header['sps_id']

        sps = sps_helper.get_sps_by_id(sps_id)
        qp, ec_part, reset_feature_memory, bit_stream = read_ip_remaining(input_buff)

        if header['nal_type'] == NalType.NAL_I:
            decoded = i_net.decompress(bit_stream, sps, qp, ec_part)
            if not args['force_intra']:
                p_net.clear_dpb()
                p_net.add_ref_feature_from_frame(decoded['x_hat'], apply_feature_adaptor=False)
        elif header['nal_type'] == NalType.NAL_P:
            decoded = p_net.decompress(bit_stream, sps, qp, ec_part, reset_feature_memory)

        recon_frame = decoded['x_hat']

        end_event.record()
        torch.cuda.synchronize(device=device)

        frame_time = start_event.elapsed_time(end_event) / 1000.0
        decoding_time.append(frame_time)

        is_intra = header['nal_type'] == NalType.NAL_I
        maximum_read = min(g_frame_delay, frame_num - decoded_frame_number)
        if is_intra:
            maximum_read = 1
        y, u, v, rgb = get_src_frame(args, src_reader, device, maximum_read, is_intra, np_only=True)

        for i in range(maximum_read):
            if isinstance(recon_frame, list):
                x_hat = recon_frame[i]
            else:
                x_hat = recon_frame
            x_hat = x_hat[:, :, :pic_height, :pic_width]
            curr_psnr, curr_ssim = get_distortion(args, x_hat, y[i], u[i], v[i], rgb[i])
            psnrs.append(curr_psnr)
            msssims.append(curr_ssim)

        if verbose >= 2:
            stream_length = 0 if bit_stream is None else len(bit_stream) * 8
            print(f'frame {decoded_frame_number} decoded, {frame_time * 1000:.3f} ms, '
                  f'bits: {stream_length}, ', end='')
            for i in range(-maximum_read, 0, 1):
                print(f'PSNR: {psnrs[i][0]:.4f}, ', end='')
            print()

        if save_decoded_frame:
            for i in range(maximum_read):
                if isinstance(recon_frame, list):
                    x_hat = recon_frame[i]
                else:
                    x_hat = recon_frame
                x_hat = x_hat[:, :, :pic_height, :pic_width]
                if args['src_type'] == 'yuv420':
                    y_rec, uv_rec = yuv_444_to_420(x_hat + 0.5)
                    y_rec = torch.clamp(y_rec * 255, 0, 255).round().byte()
                    y_rec = y_rec.squeeze(0).cpu().numpy()
                    uv_rec = torch.clamp(uv_rec * 255, 0, 255).byte()
                    uv_rec = uv_rec.squeeze(0).cpu().numpy()
                    recon_writer.write_one_frame(y_rec, uv_rec)
                else:
                    assert args['src_type'] == 'png'
                    rgb_rec = ycbcr2rgb(x_hat + 0.5)
                    rgb_rec = torch.clamp(rgb_rec * 255, 0, 255).round().byte()
                    rgb_rec = rgb_rec.squeeze(0).cpu().numpy()
                    recon_writer.write_one_frame(rgb_rec)
        decoded_frame_number += maximum_read
    input_buff.close()
    src_reader.close()

    if save_decoded_frame:
        recon_writer.close()

    test_time = time.time() - start_time
    test_time_frame_numuber = len(encoding_time)
    time_bypass_frame_num = 4  # bypass the first 4 * g_frame_delay frames as warmup
    if verbose >= 1 and test_time_frame_numuber > time_bypass_frame_num:
        encoding_time = encoding_time[time_bypass_frame_num:]
        decoding_time = decoding_time[time_bypass_frame_num:]
        avg_encoding_time = sum(encoding_time)/len(encoding_time)
        avg_decoding_time = sum(decoding_time)/len(decoding_time)
        print(f'encoding/decoding {len(encoding_time)} * {g_frame_delay} frames, '
              f'average encoding time {avg_encoding_time * 1000:.3f} ms, '
              f'average decoding time {avg_decoding_time * 1000:.3f} ms.')
    else:
        avg_encoding_time = None
        avg_decoding_time = None

    log_result = generate_log_json(frame_num, pic_height * pic_width, test_time,
                                   frame_types, bits, psnrs, msssims, verbose=verbose_json,
                                   avg_encoding_time=avg_encoding_time,
                                   avg_decoding_time=avg_decoding_time,)
    with open(args['curr_json_path'], 'w') as fp:
        json.dump(log_result, fp, indent=2)
    return log_result


g_i_net = None  # the model is initialized after each process is spawn, thus OK for multiprocess
g_p_net = None
g_frame_delay = None  # depending on the model_structure (1 or 8)


def init_func(args, gpu_num):
    set_torch_env()

    process_name = multiprocessing.current_process().name
    process_idx = int(process_name[process_name.rfind('-') + 1:])
    gpu_id = -1
    if gpu_num > 0:
        gpu_id = process_idx % gpu_num
    if gpu_id >= 0:
        if args.cuda_idx is not None:
            gpu_id = args.cuda_idx[gpu_id]
        os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
        device = 'cuda:0'
    else:
        device = 'cpu'

    # one-time setup of default cuda stream. cuda graph does not support stream 0
    custom_stream = torch.cuda.Stream(device, 0)
    torch.cuda.set_stream(custom_stream)

    model_structure = ModelStructure(args.model_structure)

    if model_structure == ModelStructure.LD:
        from src.models.video_model_ld import DMC, g_frame_delay as frame_delay
    else:
        from src.models.video_model_ht import DMC, g_frame_delay as frame_delay

    global g_i_net
    g_i_net = DMCI()
    g_i_net = g_i_net.eval()
    i_state_dict = get_state_dict(args.model_path_i)
    g_i_net.load_state_dict(i_state_dict)
    g_i_net.update(args.skip_thres)
    g_i_net = finalize_model(g_i_net, device)

    if not args.force_intra:
        global g_p_net
        if model_structure == ModelStructure.LD:
            g_p_net = DMC()
        else:
            g_p_net = DMC(model_structure=model_structure)
        g_p_net = g_p_net.eval()
        p_state_dict = get_state_dict(args.model_path_p)
        g_p_net.load_state_dict(p_state_dict)
        g_p_net.update(args.skip_thres)
        g_p_net = finalize_model(g_p_net, device)

    global g_frame_delay
    g_frame_delay = frame_delay


def worker(args):
    sub_dir_name = args['seq']
    bin_folder = os.path.join(args['stream_path'], args['ds_name'])
    create_folder(bin_folder, True)

    args['src_path'] = os.path.join(args['dataset_path'], sub_dir_name)
    args['bin_folder'] = bin_folder
    args['curr_bin_path'] = os.path.join(bin_folder, f'{args['seq']}_q{args['qp_i']}.bin')
    args['curr_rec_path'] = args['curr_bin_path'].replace('.bin', '.mp4')
    args['curr_json_path'] = args['curr_bin_path'].replace('.bin', '.json')

    result = run_one_point_with_stream(g_p_net, g_i_net, args)

    result['ds_name'] = args['ds_name']
    result['seq'] = args['seq']
    result['rate_idx'] = args['rate_idx']
    return result


def main():
    begin_time = time.time()

    args = parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError('test_video.py requires CUDA (torch.cuda.is_available() is False)')

    args.skip_thres = max(0, args.skip_thres)
    if args.cuda_idx is not None:
        cuda_device = ','.join([str(s) for s in args.cuda_idx])
        os.environ['CUDA_VISIBLE_DEVICES'] = cuda_device

    worker_num = args.worker
    assert worker_num >= 1

    with open(args.test_config) as f:
        config = json.load(f)

    gpu_num = torch.cuda.device_count()
    multiprocessing.set_start_method('spawn')
    threadpool_executor = concurrent.futures.ProcessPoolExecutor(max_workers=worker_num,
                                                                 initializer=init_func,
                                                                 initargs=(args, gpu_num))
    objs = []

    count_frames = 0
    count_sequences = 0

    rate_num = args.rate_num
    qp_i = []
    if args.qp_i is not None:
        assert len(args.qp_i) == rate_num
        qp_i = args.qp_i
    else:
        assert 2 <= rate_num <= DMCI.qp_num()
        for i in np.linspace(0, DMCI.qp_num() - 1, num=rate_num):
            qp_i.append(int(i+0.5))

    if not args.force_intra:
        if args.qp_p is not None:
            assert len(args.qp_p) == rate_num
            qp_p = args.qp_p
        else:
            qp_p = qp_i

    print(f'testing {rate_num} rates, using qp: {', '.join(str(q) for q in qp_i)}')

    root_path = args.force_root_path if args.force_root_path is not None else config['root_path']
    config = config['test_classes']
    for ds_name in config:
        if config[ds_name]['test'] == 0:
            continue
        for seq in config[ds_name]['sequences']:
            count_sequences += 1
            for rate_idx in range(rate_num):
                cur_args = {}
                cur_args['rate_idx'] = rate_idx
                cur_args['qp_i'] = qp_i[rate_idx]
                if not args.force_intra:
                    cur_args['qp_p'] = qp_p[rate_idx]
                cur_args['force_intra'] = args.force_intra
                cur_args['reset_interval'] = args.reset_interval
                cur_args['seq'] = seq
                cur_args['src_type'] = config[ds_name]['src_type']
                cur_args['src_height'] = config[ds_name]['sequences'][seq]['height']
                cur_args['src_width'] = config[ds_name]['sequences'][seq]['width']
                cur_args['intra_period'] = config[ds_name]['sequences'][seq]['intra_period']
                if args.force_intra:
                    cur_args['intra_period'] = 1
                if args.force_intra_period > 0:
                    cur_args['intra_period'] = args.force_intra_period
                cur_args['frame_num'] = config[ds_name]['sequences'][seq]['frames']
                if args.force_frame_num > 0:
                    cur_args['frame_num'] = args.force_frame_num
                cur_args['calc_ssim'] = args.calc_ssim
                cur_args['dataset_path'] = os.path.join(root_path, config[ds_name]['base_path'])
                cur_args['check_existing'] = args.check_existing
                cur_args['stream_path'] = args.stream_path
                cur_args['save_decoded_frame'] = args.save_decoded_frame
                cur_args['ds_name'] = ds_name
                cur_args['verbose'] = args.verbose
                cur_args['verbose_json'] = args.verbose_json

                count_frames += cur_args['frame_num']

                obj = threadpool_executor.submit(worker, cur_args)
                objs.append(obj)

    results = []
    for obj in tqdm(objs):
        result = obj.result()
        results.append(result)

    log_result = {}
    for ds_name in config:
        if config[ds_name]['test'] == 0:
            continue
        log_result[ds_name] = {}
        for seq in config[ds_name]['sequences']:
            log_result[ds_name][seq] = {}

    for res in results:
        log_result[res['ds_name']][res['seq']][f'{res['rate_idx']:03d}'] = res

    out_json_dir = os.path.dirname(args.output_path)
    if len(out_json_dir) > 0:
        create_folder(out_json_dir, True)
    with open(args.output_path, 'w') as fp:
        dump_json(log_result, fp, float_digits=6, indent=2)

    total_minutes = (time.time() - begin_time) / 60
    print('Test finished')
    print(f'Tested {count_frames} frames from {count_sequences} sequences')
    print(f'Total elapsed time: {total_minutes:.1f} min')


if __name__ == '__main__':
    main()