Use FlowNet1 and FlowNet2 to predict separately

This article has no code installation and environment configuration section

FlowNet1

The images are placed imageunder the folder:

The pre-trained model is placed on the Baidu network disk , extract the code: 2022
run the command :

python run_inference.py ./images ./pretrained/flownets_EPE1.951.pth.tar

FlowNet2

run pngimage

The predicted image adopts the pngformat

of pre-training model address, Baidu network disk , extraction code:2022

import torch
import numpy as np
import argparse
import time
from models import FlowNet2  # the path is depended on where you create this module

from utils.flow_utils import flow2img
import matplotlib.pyplot as plt
import cv2
image_path = './image/'
save_path = './image/result/'	#结果保存位置
image_name = 'image1_' #图像名称，更改序号来实现预测不同的图像对
model_path = "/home/ly/code/flownet2/models/checkpoints/FlowNet2_checkpoint.pth.tar" # 模型位置

'''
crop_size 需要根据图像的实际尺寸更改，且是32的倍数
'''
crop_size = (1024, 1024) 
prvs = cv2.imread(image_path + image_name+'1.png')
next = cv2.imread(image_path + image_name+'2.png')
#pim1 = prvs
#pim2 = next
pim1 = cv2.resize(prvs, crop_size, interpolation = cv2.INTER_CUBIC)
pim2 = cv2.resize(next, crop_size, interpolation = cv2.INTER_CUBIC)

parser = argparse.ArgumentParser()
parser.add_argument('--fp16', action='store_true', help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).')
parser.add_argument("--rgb_max", type=float, default=255.)

args = parser.parse_args()
net = FlowNet2(args).cuda()
dict = torch.load(model_path)
net.load_state_dict(dict["state_dict"])

images = [pim1, pim2]
images = np.array(images).transpose(3, 0, 1, 2)
im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()

start = time.time()
result = net(im).squeeze()
end = time.time()
print(end-start)
data = result.data.cpu().numpy().transpose(1, 2, 0)
img = flow2img(data)
cv2.imwrite(save_path + image_name[:-1] +'.png',img)

run video file

# -*- coding: utf-8 -*-
import numpy as np
import cv2
import warnings
warnings.filterwarnings("ignore")
import torch
import numpy as np
import argparse
import time
from models import FlowNet2  # the path is depended on where you create this module

from utils.flow_utils import flow2img
import matplotlib.pyplot as plt
# 输入视频地址
cap = cv2.VideoCapture('./test.mp4')

#获取第一帧
ret, frame1 = cap.read()
prvs = frame1
i = 0 #控制实现的张数
save_path = './result/'
while(1):
    ret, frame2 = cap.read()

    next = frame2

    crop_size = (1024, 384)
    pim1 = cv2.resize(prvs, crop_size, interpolation = cv2.INTER_CUBIC)
    pim2 = cv2.resize(next, crop_size, interpolation = cv2.INTER_CUBIC)

    # obtain the necessary args for construct the flownet framework
    parser = argparse.ArgumentParser()
    parser.add_argument('--fp16', action='store_true', help='Run model in pseudo-fp16 mode (fp16 storage fp32 math).')
    parser.add_argument("--rgb_max", type=float, default=255.)

    args = parser.parse_args()

    # initial a Net
    net = FlowNet2(args).cuda()
    # 加载模型
    dict = torch.load("/home/ly/code/flownet2/models/checkpoints/FlowNet2_checkpoint.pth.tar")
    net.load_state_dict(dict["state_dict"])

    images = [pim1, pim2]
    images = np.array(images).transpose(3, 0, 1, 2)
    im = torch.from_numpy(images.astype(np.float32)).unsqueeze(0).cuda()

    start = time.time()
    result = net(im).squeeze()
    end = time.time()
    print(end-start)
    data = result.data.cpu().numpy().transpose(1, 2, 0)
    img = flow2img(data)
    cv2.imwrite(save_path + str(i)+'.png',img)
    i = i+1
    prvs = next

PNG image to video (MP4 file)

import os  # python标准库，不需要安装，用于系统文件操作相关
import cv2  # python非标准库，pip install opencv-python 多媒体处理
from PIL import Image  # python非标准库，pip install pillow，图像处理
import moviepy.editor as mov  # python非标准库，pip install moviepy，多媒体编辑

def image_to_video(image_path, media_path):
    '''
    图片合成视频函数
    :param image_path: 图片路径
    :param media_path: 合成视频保存路径
    :return:
    '''
    # 获取图片路径下面的所有图片名称
    image_names = os.listdir(image_path)
    print(image_names)
    # 对提取到的图片名称进行排序
    image_names.sort(key=lambda n: int(n[:-4]))
    # 设置写入格式
    fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')
    # 设置每秒帧数
    fps = 10  # 由于图片数目较少，这里设置的帧数比较低
    # 读取第一个图片获取大小尺寸，因为需要转换成视频的图片大小尺寸是一样的
    image = Image.open(image_path + image_names[0])
    # 初始化媒体写入对象
    media_writer = cv2.VideoWriter(media_path, fourcc, fps, image.size)
    # 遍历图片，将每张图片加入视频当中
    for image_name in image_names:
        im = cv2.imread(os.path.join(image_path, image_name))
        media_writer.write(im)
        print(image_name, '合并完成！')
    # 释放媒体写入对象
    media_writer.release()
    print('无声视频写入完成！')
# 输入存放图像的文件夹路径，输出视频的文件
image_to_video('./result/', './result.mp4')