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()