demo_video.ipynb

from __future__ import print_function, division
import matplotlib.pyplot as plt
import os
import argparse
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from model.cnn_geometric_model import CNNGeometric
from data.pf_dataset import PFDataset
from data.download_datasets import download_PF_willow
from image.normalization import NormalizeImageDict, normalize_image
from util.torch_util import BatchTensorToVars, str_to_bool
from geotnf.transformation import GeometricTnf
from geotnf.point_tnf import *
import matplotlib.pyplot as plt
from skimage import io
import warnings
from torchvision.transforms import Normalize
from collections import OrderedDict
import numpy as np
from IPython.display import clear_output

def visualize(image):
#     plt.figure(figsize=(4, 4))
    plt.axis('off')
    plt.imshow(image)
    plt.show()
    
resizeCNN = GeometricTnf(out_h=240, out_w=240, use_cuda = False) 

def preprocess_image(image):
    # convert to torch Variable
    image = np.expand_dims(image.transpose((2,0,1)),0)
    image = torch.Tensor(image.astype(np.float32)/255.0)
    image_var = Variable(image,requires_grad=False)

    # Resize image using bilinear sampling with identity affine tnf
    image_var = resizeCNN(image_var)
    
    # Normalize image
    image_var = normalize_image(image_var)
    
    return image_var
    
    
use_cuda = torch.cuda.is_available()
def build_model():
    feature_extraction_cnn = 'vgg'
    
    print(use_cuda)

#     model_aff_path = '/home/41875/zlk_236150/AIstation_back/best_checkpoint_adam_offset_grid_lossvgg.pth.tar'
    
    # server best epoch 108:  validation average loss:0.0001   train_loss:0.00006
    model_aff_path = '/home/41875/zlk_236150/openProject/cnngeometric_pytorch/trained_models/best_checkpoint_adam_offset_grid_lossvgg.pth.tar' 

    # Create model
    print('Creating CNN model...')

    model_aff = CNNGeometric(use_cuda=use_cuda,feature_extraction_cnn=feature_extraction_cnn,output_dim=2)

    checkpoint = torch.load(model_aff_path, map_location=lambda storage, loc: storage)
    checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
    model_aff.load_state_dict(checkpoint['state_dict'])
    
    model_aff.eval()
    print("create model succeed")
    
    return model_aff

affine_matrix_p1 = torch.tensor([[1.0, 0.0],[0.0, 1.0]])
affine_matrix_p1 = affine_matrix_p1.cuda()
affine_matrix_p1 = affine_matrix_p1.repeat(1,1,1)

def registerImages(model_aff,source_image,target_image):

    source_image_var = preprocess_image(source_image)
    target_image_var = preprocess_image(target_image)
    if use_cuda:
        source_image_var = source_image_var.cuda()
        target_image_var = target_image_var.cuda()
        
    batch = {'source_image': source_image_var, 'target_image':target_image_var}
    
    theta_aff = model_aff(batch)
    theta_aff = theta_aff.unsqueeze(2)
    theta_aff = torch.cat((affine_matrix_p1,theta_aff),2)
    
    return theta_aff,batch

affTnf = GeometricTnf(geometric_model='offset', use_cuda=use_cuda)
normalizeTnf = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
resizeTgt = GeometricTnf(out_h=240, out_w=240, use_cuda = use_cuda) 

def aff_images(batchs,theta_aff):
    warped_image_aff = affTnf(batch['source_image'],theta_aff.view(-1,2,3))
    warped_image_aff_np = normalize_image(resizeTgt(warped_image_aff),forward=False).data.squeeze(0).transpose(0,1).transpose(1,2).cpu().numpy()
    return warped_image_aff_np

model_aff = build_model()


def show_compares(source_image,target_image,warped_image_aff_np):
    source_image = np.clip(source_image,a_min=0,a_max=1)
    target_image = np.clip(target_image,a_min=0,a_max=1)
    warped_image_aff_np = np.clip(warped_image_aff_np,a_min=0,a_max=1)
    N_subplots = 3
    fig, axs = plt.subplots(1,N_subplots)
    axs[0].imshow(source_image)
    axs[0].set_title('origin')
    axs[1].imshow(target_image)
    axs[1].set_title('target')
    subplot_idx = 2
    
    axs[subplot_idx].imshow(warped_image_aff_np)
    axs[subplot_idx].set_title('aff')
    subplot_idx +=1 

    for i in range(N_subplots):
        axs[i].axis('off')
        
    fig.set_dpi(150)
    plt.show()
    
def plot_xy_curve(theta_xy):
    plt.figure()
#     print(theta_xy.shape)
    offset_x = theta_xy[:,0]
    offset_y = theta_xy[:,1]
#     print(offset_x.shape)
    
    plt.plot([i for i in range(len(offset_x))], offset_x,label="offset_x")
    plt.plot([i for i in range(len(offset_x))], offset_y,label="offset_y")

    plt.title('offset-frame')
    plt.xlabel('frame')
    plt.ylabel('offset')

    plt.show()

import time
import cv2
warnings.filterwarnings("ignore")
current_time = 0
 
# 图像处理函数
def processImg(img):
    # 画出一个框
    #cv2.rectangle(img, (500, 300), (800, 400), (0, 0, 255), 5, 1, 0)
    # 上下翻转
    # img= cv2.flip(img, 0)

    # 显示FPS
    global current_time
    if current_time == 0:
        current_time = time.time()
    else:
        last_time = current_time
        current_time = time.time()
        fps = 1. / (current_time - last_time)
        text = "FPS: %d" % int(fps)
        cv2.putText(img, text , (0,100), cv2.FONT_HERSHEY_TRIPLEX, 3.65, (255, 0, 0), 2)

    return img
 