gpucuda

import os.path
import logging
import time
import numpy as np
import math
import sys
sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')

import cv2
from collections import OrderedDict

import torch

from utils import utils_logger
from utils import utils_image as util

from models.network_ffdnet import FFDNet as net
from subprocess import call
import line_profiler


def main():


    start_2 = time.time()
    model_name = 'ffdnet_color'           
    task_current = 'dn'       
    sf = 1                    
    if 'color' in model_name:
        n_channels = 3        
        nc = 96               
        nb = 12              
    else:
        n_channels = 1        
        nc = 64               
        nb = 15             

    model_pool = 'model'  # fixed

    border = sf if task_current == 'sr' else 0     # shave boader to calculate PSNR and SSIM
    model_path = os.path.join(model_pool, model_name+'.pth')

    #device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    device = torch.device('cpu')
    model = net(in_nc=n_channels, out_nc=n_channels, nc=nc, nb=nb, act_mode='R')
    model.load_state_dict(torch.load(model_path), strict=True)
    model.eval()
    
    for k, v in model.named_parameters():
        v.requires_grad = False
    model = model.to(device)
    
    cap = cv2.VideoCapture('video/MOV_0005.MOV')
    cap_num = cap.get(7)
    cap_width = math.ceil(cap.get(3))
    cap_height = math.ceil(cap.get(4))
    cap_fps = math.ceil(cap.get(5))

    size = (cap_width, cap_height)
    #fourcc = cv2.VideoWriter_fourcc('I','4','2','0')
    fourcc = cv2.VideoWriter_fourcc('m','p','4','v')
    video = cv2.VideoWriter('results/result.mp4',fourcc, cap_fps, size)

    start_3 = time.time()
    suc, frame = cap.read()  
    cnt = 1  

    while suc:       
    #for i in range(40):   
           
        start_2 = time.time() 
        img_L = util.uint2tensor4(frame)
        img_L = img_L.to(device)
        
        sigma = torch.full((1,1,1,1), 0.098).type_as(img_L)
        img_E = model(img_L, sigma)
        end_2 = time.time()
        print('time 2:' + str(end_2-start_2))

        start_2 = time.time()        
        img_E = util.tensor2uint(img_E)
        #if(cnt%5==0):
        end_2 = time.time()
        print('time 4:' + str(end_2-start_2))
        
        
        #img_L = util.tensor2uint(img_L)
        #util.imshow(np.concatenate([img_E, img_L], axis=1), title='Recovered / Ground-truth')
        
        video.write(img_E)
        suc, frame = cap.read()
        cnt = cnt+1
        
    end_3 = time.time()
    print('time 3:' + str(end_3-start_3))
    cap.release()
    video.release()
    

if __name__ == '__main__':

    #profile = line_profiler.LineProfiler(main)
    #profile.enable()
    main()
    #profile.disable()
    #profile.print_stats()


 

你可能感兴趣的:(gpucuda)