diff --git a/test.py b/test.py index 3f67e07..408e074 100644 --- a/test.py +++ b/test.py @@ -53,18 +53,26 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c total_time = 0 I_last_lab_predict = None - + + IB_lab_large = transform(frame_ref).unsqueeze(0).cuda() + IB_lab = torch.nn.functional.interpolate(IB_lab_large, scale_factor=0.5, mode="bilinear") + IB_l = IB_lab[:, 0:1, :, :] + IB_ab = IB_lab[:, 1:3, :, :] + with torch.no_grad(): + I_reference_lab = IB_lab + I_reference_l = I_reference_lab[:, 0:1, :, :] + I_reference_ab = I_reference_lab[:, 1:3, :, :] + I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) + features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) + for index, frame_name in enumerate(tqdm(filenames)): frame1 = Image.open(os.path.join(input_path, frame_name)) IA_lab_large = transform(frame1).unsqueeze(0).cuda() - IB_lab_large = transform(frame_ref).unsqueeze(0).cuda() IA_lab = torch.nn.functional.interpolate(IA_lab_large, scale_factor=0.5, mode="bilinear") - IB_lab = torch.nn.functional.interpolate(IB_lab_large, scale_factor=0.5, mode="bilinear") IA_l = IA_lab[:, 0:1, :, :] IA_ab = IA_lab[:, 1:3, :, :] - IB_l = IB_lab[:, 0:1, :, :] - IB_ab = IB_lab[:, 1:3, :, :] + if I_last_lab_predict is None: if opt.frame_propagate: I_last_lab_predict = IB_lab @@ -74,15 +82,6 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c # start the frame colorization with torch.no_grad(): I_current_lab = IA_lab - I_reference_lab = IB_lab - I_reference_l = I_reference_lab[:, 0:1, :, :] - I_reference_ab = I_reference_lab[:, 1:3, :, :] - I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)) - - # t_start = time.clock() - # torch.cuda.synchronize() - - features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True) I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization( I_current_lab, I_reference_lab, @@ -96,14 +95,6 @@ def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, c ) I_last_lab_predict = torch.cat((IA_l, I_current_ab_predict), dim=1) - # update timing - # torch.cuda.synchronize() - # delta_t = time.clock() - t_start - # print("runtime:", delta_t) - # if iter_num > 0: - # total_time += delta_t - # print("%.2g second" % (total_time / iter_num)) - # upsampling curr_bs_l = IA_lab_large[:, 0:1, :, :] curr_predict = (