inference_video_w.py

import os
import cv2
import torch
import numpy as np
from tqdm import tqdm
from torch.nn import functional as F
import warnings
import _thread
import skvideo.io
from queue import Queue, Empty
from model.pytorch_msssim import ssim_matlab
import shutil
import tempfile
import time

warnings.filterwarnings("ignore")

# Utility class to mimic argparse object
class Args:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

def transferAudio(sourceVideo, targetVideo):
    # generate a unique temp directory for this user
    unique_temp_dir = tempfile.mkdtemp()
    tempAudioFileName = os.path.join(unique_temp_dir, "audio.mkv")

    # extract audio from video
    os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -c:a copy -vn {}'.format(sourceVideo, tempAudioFileName))

    targetNoAudio = os.path.splitext(targetVideo)[0] + "_noaudio" + os.path.splitext(targetVideo)[1]
    os.rename(targetVideo, targetNoAudio)
    # combine audio file and new video file
    os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))

    if os.path.getsize(targetVideo) == 0: # if ffmpeg failed to merge the video and audio together try converting the audio to aac
        tempAudioFileName = os.path.join(unique_temp_dir, "audio.m4a")
        os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -c:a aac -b:a 160k -vn {}'.format(sourceVideo, tempAudioFileName))
        os.system('ffmpeg -hide_banner -loglevel error -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
        if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format
            os.rename(targetNoAudio, targetVideo)
            print("Audio transfer failed. Interpolated video will have no audio")
        else:
            print("Lossless audio transfer failed. Audio was transcoded to AAC (M4A) instead.")
            # remove audio-less video
            os.remove(targetNoAudio)
    else:
        os.remove(targetNoAudio)

    # remove temp directory
    shutil.rmtree(unique_temp_dir)

def inference(
    video=None,
    output=None,
    img=None,
    montage=False,
    modelDir='train_log',
    fp16=False,
    UHD=False,
    scale=1.0,
    skip=False,
    fps=None,
    png=False,
    ext='mp4',
    exp=1,
    multi=2
):
    # Initialize Arguments Object
    args = Args(
        video=video, output=output, img=img, montage=montage,
        modelDir=modelDir, fp16=fp16, UHD=UHD, scale=scale,
        skip=skip, fps=fps, png=png, ext=ext, exp=exp, multi=multi
    )

    # Argument Logic Adjustment
    if args.exp != 1:
        args.multi = (2 ** args.exp)
    
    # Assertions
    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

    # Device Setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.set_grad_enabled(False)
    if torch.cuda.is_available():
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        if(args.fp16):
            torch.set_default_tensor_type(torch.cuda.HalfTensor)

    # Load Model
    from train_log.RIFE_HDv3 import Model
    model = Model()
    if not hasattr(model, 'version'):
        model.version = 0
    model.load_model(args.modelDir, -1)
    print("Loaded 3.x/4.x HD model.")
    model.eval()
    model.device()

    # Video/Image Setup
    if not args.video is None:
        videoCapture = cv2.VideoCapture(args.video)
        original_fps = videoCapture.get(cv2.CAP_PROP_FPS)
        tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
        videoCapture.release()
        
        if args.fps is None or args.fps == 0:
            fpsNotAssigned = True
            args.fps = original_fps * args.multi
        else:
            fpsNotAssigned = False
            
        videogen = skvideo.io.vreader(args.video)
        lastframe = next(videogen)
        # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') # Unused in original logic for skvideo
        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, original_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!")
    else:
        videogen = []
        for f in os.listdir(args.img):
            if 'png' in f:
                videogen.append(f)
        tot_frame = len(videogen)
        videogen.sort(key= lambda x:int(x[:-4]))
        lastframe = cv2.imread(os.path.join(args.img, videogen[0]), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
        videogen = videogen[1:]
    
    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)
        
        outputdict = {
            '-c:v': 'libx264', 
            '-crf': '17', 
            '-preset': 'slow',
            '-pix_fmt': 'yuv420p'
        }
        vid_out = skvideo.io.FFmpegWriter(vid_out_name, inputdict={'-r': str(args.fps)}, outputdict=outputdict)

    # --- Nested Helper Functions to capture 'args', 'model', 'vid_out' scope ---

    def clear_write_buffer(write_buffer):
        cnt = 0
        while True:
            item = write_buffer.get()
            if item is None:
                break
            if args.png:
                cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1])
                cnt += 1
            else:
                vid_out.writeFrame(item)

    def build_read_buffer(read_buffer, videogen):
        try:
            for frame in videogen:
                if not args.img is None:
                    frame = cv2.imread(os.path.join(args.img, frame), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
                if args.montage:
                    frame = frame[:, left: left + w]
                read_buffer.put(frame)
        except:
            pass
        read_buffer.put(None)

    def make_inference(I0, I1, n):    
        if model.version >= 3.9:
            res = []
            for i in range(n):
                res.append(model.inference(I0, I1, (i+1) * 1. / (n+1), args.scale))
            return res
        else:
            middle = model.inference(I0, I1, args.scale)
            if n == 1:
                return [middle]
            first_half = make_inference(I0, middle, n=n//2)
            second_half = make_inference(middle, I1, n=n//2)
            if n%2:
                return [*first_half, middle, *second_half]
            else:
                return [*first_half, *second_half]

    def pad_image(img):
        if(args.fp16):
            return F.pad(img, padding).half()
        else:
            return F.pad(img, padding)

    # --- Pre-Loop Setup ---

    left = 0 # Define default
    if args.montage:
        left = w // 4
        w = w // 2
    
    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)
    
    pbar = tqdm(total=tot_frame)
    if args.montage:
        lastframe = lastframe[:, left: left + w]
        
    write_buffer = Queue(maxsize=500)
    read_buffer = Queue(maxsize=500)
    
    # Start threads
    _thread.start_new_thread(build_read_buffer, (read_buffer, videogen))
    _thread.start_new_thread(clear_write_buffer, (write_buffer,))

    I1 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
    I1 = pad_image(I1)
    temp = None 

    # --- Main Loop ---
    
    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 = pad_image(I1)
        I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False)
        I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
        ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])

        break_flag = False
        if ssim > 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 = pad_image(I1)
            I1 = model.inference(I0, I1, scale=args.scale)
            I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
            ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
            frame = (I1[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w]
            
        if ssim < 0.2:
            output_frames = []
            for i in range(args.multi - 1):
                output_frames.append(I0)
        else:
            output_frames = make_inference(I0, I1, args.multi - 1)

        if args.montage:
            write_buffer.put(np.concatenate((lastframe, lastframe), 1))
            for mid in output_frames:
                mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
                write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
        else:
            write_buffer.put(lastframe)
            for mid in output_frames:
                mid = (((mid[0] * 255.).byte().cpu().numpy().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.close()

    # Audio Transfer Logic
    if args.png == False and fpsNotAssigned == True and not args.video is None:
        try:
            transferAudio(args.video, vid_out_name)
        except:
            print("Audio transfer failed. Interpolated video will have no audio")
            targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1]
            os.rename(targetNoAudio, vid_out_name)
            
    return vid_out_name