run_axmodel.py

import os
import cv2
import time
import argparse
import numpy as np
import axengine as axe
import _thread
import torch
import torch.nn.functional as F
import ms_ssim

from tqdm import tqdm
from queue import Queue, Empty

parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
parser.add_argument('--video', dest='video', type=str, default='./demo.mp4')
parser.add_argument('--output', dest='output', type=str, default=None)
parser.add_argument('--img', dest='img', type=str, default=None)
parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video')
parser.add_argument('--model', dest='model', type=str, default=None, help='directory with trained model files')
parser.add_argument('--fp16', dest='fp16', action='store_true', help='fp16 mode for faster and more lightweight inference on cards with Tensor Cores')
parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video')
parser.add_argument('--scale', dest='scale', type=float, default=1.0, help='Try scale=0.5 for 4k video')
parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing')
parser.add_argument('--fps', dest='fps', type=int, default=None)
parser.add_argument('--png', dest='png', action='store_true', help='whether to vid_out png format vid_outs')
parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='vid_out video extension')
parser.add_argument('--exp', dest='exp', type=int, default=1)
parser.add_argument('--multi', dest='multi', type=int, default=2)

def read_video(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise IOError(f"Cannot open video: {video_path}")
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  
            yield frame
    finally:
        cap.release()
    
def clear_write_buffer(user_args, write_buffer, vid_out):
    cnt = 0
    while True:
        item = write_buffer.get()
        if item is None:
            break
        if user_args.png:
            cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1])
            cnt += 1
        else:
            vid_out.write(item[:, :, ::-1])

def build_read_buffer(user_args, read_buffer, videogen):
    try:
        for frame in videogen:
            if not user_args.img is None:
                frame = cv2.imread(os.path.join(user_args.img, frame), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
            if user_args.montage:
                frame = frame[:, left: left + w]
            read_buffer.put(frame)
    except:
        pass
    read_buffer.put(None)
    
def pad_image(img, padding):
    if(args.fp16):
        return F.pad(img, padding).half()
    else:
        return F.pad(img, padding)
    
def run(args):
    '''onnx inference'''
    # model
    session = axe.InferenceSession(args.model)
    output_names = [x.name for x in session.get_outputs()]
    input_name = session.get_inputs()[0].name 
    
    # video
    videoCapture = cv2.VideoCapture(args.video)
    fps = videoCapture.get(cv2.CAP_PROP_FPS)
    tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
    videoCapture.release()
    if args.fps is None:
        fpsNotAssigned = True
        args.fps = fps * args.multi
    else:
        fpsNotAssigned = False
    videogen = read_video(args.video)
    lastframe = next(videogen)
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
    video_path_wo_ext, ext = os.path.splitext(args.video)
    print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args.ext, tot_frame, fps, args.fps))
    if args.png == False and fpsNotAssigned == True:
        print("The audio will be merged after interpolation process")
    else:
        print("Will not merge audio because using png or fps flag!")

    # 
    h, w, _ = lastframe.shape
    vid_out_name = None
    vid_out = None
    if args.png:
        if not os.path.exists('vid_out'):
            os.mkdir('vid_out')
    else:
        if args.output is not None:
            vid_out_name = args.output
        else:
            vid_out_name = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, args.multi, int(np.round(args.fps)), args.ext)
        vid_out = cv2.VideoWriter(vid_out_name, fourcc, args.fps, (w, h))
    
    tmp = max(128, int(128 / args.scale))
    ph = ((h - 1) // tmp + 1) * tmp
    pw = ((w - 1) // tmp + 1) * tmp
    #padding = (0, pw - w, 0, ph - h)
    padding = ((0, 0), (0, 0), (0, ph - h), (0, pw - w))
    pbar = tqdm(total=tot_frame, ncols=80)
    
    write_buffer = Queue(maxsize=500)
    read_buffer = Queue(maxsize=500)
    _thread.start_new_thread(build_read_buffer, (args, read_buffer, videogen))
    _thread.start_new_thread(clear_write_buffer, (args, write_buffer, vid_out))

    #device = 'cpu'
    #I1 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
    I1 = np.expand_dims(np.transpose(lastframe, (2,0,1)), 0).astype(np.float32) / 255.
    I1 = np.pad(I1, padding)
    
    temp = None # save lastframe when processing static frame
    while True:
        if temp is not None:
            frame = temp
            temp = None
        else:
            frame = read_buffer.get()
        if frame is None:
            break
        I0 = I1
        #I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
        I1 = np.expand_dims(np.transpose(frame, (2,0,1)), 0).astype(np.float32) / 255.
        I1 = np.pad(I1, padding)

        I0_small = F.interpolate(torch.from_numpy(I0).float(), (32, 32), mode='bilinear', align_corners=False)
        I1_small = F.interpolate(torch.from_numpy(I1).float(), (32, 32), mode='bilinear', align_corners=False)
        ssim = ms_ssim.ssim_matlab(I0_small[:, :3], I1_small[:, :3])

        break_flag = False
        if ssim > 0.996:  #0.996
            frame = read_buffer.get() # read a new frame
            if frame is None:
                break_flag = True
                frame = lastframe
            else:
                temp = frame
            #I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
            I1 = np.expand_dims(np.transpose(frame, (2,0,1)), 0).astype(np.float32) / 255.
            I1 = np.pad(I1, padding)

            #imgs = torch.cat((I0, I1), 1).cpu().numpy()
            imgs = np.concatenate((I0, I1), axis=1)
            I1 = session.run(output_names, {input_name: imgs})

            #I1 = torch.from_numpy(I1[-1])
            I1 = np.array(I1[-1])
            I1_small = F.interpolate(torch.from_numpy(I1).float(), (32, 32), mode='bilinear', align_corners=False)
            ssim = ms_ssim.ssim_matlab(I0_small[:, :3], I1_small[:, :3])
            #frame = (I1[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w]
            frame = np.clip(I1[0] * 255, 0, 255).astype(np.uint8).transpose(1, 2, 0)[:h, :w]
            
        if ssim < 0.2:
            output = []
            for i in range(args.multi - 1):
                output.append(I0)
            '''
            output = []
            step = 1 / args.multi
            alpha = 0
            for i in range(args.multi - 1):
                alpha += step
                beta = 1-alpha
                output.append(torch.from_numpy(np.transpose((cv2.addWeighted(frame[:, :, ::-1], alpha, lastframe[:, :, ::-1], beta, 0)[:, :, ::-1].copy()), (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.)
            '''
        else:
            imgs = np.concatenate((I0, I1), axis=1)
            output = [session.run(output_names, {input_name: imgs})[-1]]

        if args.montage:
            write_buffer.put(np.concatenate((lastframe, lastframe), 1))
            for mid in output:
                #mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
                mid = np.clip(mid[0] * 255, 0, 255).astype(np.uint8).transpose(1, 2, 0)
                write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
        else:
            write_buffer.put(lastframe)
            for mid in output:
                #mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
                mid = np.clip(mid[0] * 255, 0, 255).astype(np.uint8).transpose(1, 2, 0)
                write_buffer.put(mid[:h, :w])
        pbar.update(1)
        lastframe = frame
        if break_flag:
            break
        
    if args.montage:
        write_buffer.put(np.concatenate((lastframe, lastframe), 1))
    else:
        write_buffer.put(lastframe)      
    write_buffer.put(None)

    while(not write_buffer.empty()):
        time.sleep(0.1)
    pbar.close()
    if not vid_out is None:
        vid_out.release()
        
if __name__ == '__main__':
    args = parser.parse_args()
    if args.exp != 1:
        args.multi = (2 ** args.exp)
    assert (not args.video is None or not args.img is None)
    if args.skip:
        print("skip flag is abandoned, please refer to issue #207.")
    if args.UHD and args.scale==1.0:
        args.scale = 0.5
    assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0]
    if not args.img is None:
        args.png = True
        
    run(args)