fast style transfer 快速风格转换 导出视频

原文链接: fast style transfer 快速风格转换 导出视频

上一篇: conda 使用 pyinstaller 打包py程序为exe文件

下一篇: dqn 打砖块 Env 对原始gym的Env进行包装

效果

https://www.bilibili.com/video/av45157478

使用sk写的文件虽然可以播放但不能再pr中使用,需要使用cv才行...

opencv 版

import cv2 as cv
import tensorflow as tf
import matplotlib.pyplot as plt
import scipy.misc as sm

import numpy as np

import skvideo

FFmpeg_path = r"C:\Users\Ace\Downloads\ffmpeg-20190227-85051fe-win64-static\bin"
skvideo.setFFmpegPath(FFmpeg_path)

from skvideo import io

pb_path = './pb/黑白.pb'
# pb_path = './pb/mnls.pb'
# video_path = 'xjtu.mp4'
video_path = r"C:\Users\Ace\Videos\Cinemagraph of XJTU.mp4"
# video_path = r"C:\Users\Ace\Videos\4K航拍-贵州黄果树瀑布.mp4"
out_path = './video/out.mp4'

sess = tf.Session()
output_graph_def = tf.GraphDef()

with open(pb_path, "rb") as f:
    output_graph_def.ParseFromString(f.read())
    tf.import_graph_def(
        output_graph_def,
        name='',  # 默认name为import,类似scope
        # return_elements=['generator/mul:0']
    )
sess.run(tf.global_variables_initializer())
in_image = sess.graph.get_tensor_by_name("in_x:0")
in_height = sess.graph.get_tensor_by_name("in_height:0")
in_width = sess.graph.get_tensor_by_name("in_width:0")
output = sess.graph.get_tensor_by_name("generator/output:0")
# (1, ?, ?, 3) <unknown> <unknown> (?, ?, ?, ?)
print(in_image.shape, in_height.shape, in_width.shape, output.shape)

capture = cv.VideoCapture(video_path)


def get_styled(img):
    _, h, w, c = img.shape
    out_val = sess.run(
        output, {
            in_image: img,
            in_height: h,
            in_width: w,
        }
    )
    # out_val = np.squeeze(out_val, 0)
    out_val = np.clip(out_val, 0, 255).astype(np.uint8)
    return out_val


# 设置格式
fourcc = cv.VideoWriter_fourcc(*"mp4v")

frame_rate = 30
flag, frame = capture.read()
# 设置输出文件路径和大小
out_file = cv.VideoWriter(out_path, fourcc, frame_rate, frame.shape[:2])

q = []
batch_size = 1
cnt = 0
while True:
    flag, frame = capture.read()
    cnt += 1
    if cnt < 500:
        continue
    if not flag:
        break
    # cv.imshow('frame', frame)
    print(frame.shape)
    # frame = frame[:, :, ::-1]
    q.append(frame)
    if len(q) == batch_size:
        frames = np.stack(q)
        q = []
        new_imgs = get_styled(frames)
        # cv.imshow('new_img', new_img)
        # print(new_img.shape)
        for new_img in new_imgs:
            out_file.write(new_img)
            cv.imwrite('tmp.jpg',new_img)

sk

import cv2 as cv
import tensorflow as tf
import matplotlib.pyplot as plt
import scipy.misc as sm

import numpy as np

import skvideo

FFmpeg_path = r"C:\Users\Ace\Downloads\ffmpeg-20190227-85051fe-win64-static\bin"
skvideo.setFFmpegPath(FFmpeg_path)

from skvideo import io
import scipy.misc as sm
import cv2 as cv

# pb_path = './pb/黑白.pb'
pb_path = './pb/starry.pb'
# pb_path = './pb/mnls.pb'
video_path = r"C:\Users\Ace\Videos\Cinemagraph of XJTU.mp4"
# video_path = r"C:\Users\Ace\Videos\4K航拍-贵州黄果树瀑布.mp4"
out_path = './video/out2.mp4'

sess = tf.Session()
output_graph_def = tf.GraphDef()

with open(pb_path, "rb") as f:
    output_graph_def.ParseFromString(f.read())
    tf.import_graph_def(
        output_graph_def,
        name='',  # 默认name为import,类似scope
        # return_elements=['generator/mul:0']
    )
sess.run(tf.global_variables_initializer())
in_image = sess.graph.get_tensor_by_name("in_x:0")
in_height = sess.graph.get_tensor_by_name("in_height:0")
in_width = sess.graph.get_tensor_by_name("in_width:0")
output = sess.graph.get_tensor_by_name("generator/output:0")
# (1, ?, ?, 3) <unknown> <unknown> (?, ?, ?, ?)
print(in_image.shape, in_height.shape, in_width.shape, output.shape)

capture = cv.VideoCapture(video_path)


def get_styled(img):
    img = np.expand_dims(img, 0).astype(np.float32)
    _, h, w, c = img.shape
    out_val = sess.run(
        output, {
            in_image: img,
            in_height: h,
            in_width: w,
        }
    )
    out_val = np.squeeze(out_val, 0)
    out_val = np.clip(out_val, 0, 255).astype(np.uint8)
    return out_val


reader = skvideo.io.vreader(video_path)

writer = skvideo.io.FFmpegWriter(out_path, outputdict={
    # '-f': 'mp4', '-b': '300000000'
    '-vcodec': 'libx264', '-r': '25'
})
fourcc = cv.VideoWriter_fourcc(*"mp4v")

# 设置输出文件路径和大小
out_file = cv.VideoWriter("out_cv.mp4", fourcc, 25, (800, 600))

for frame in reader:
    # print(frame.shape)
    frame = sm.imresize(frame, (600, 800))
    out_img = get_styled(frame)
    writer.writeFrame(out_img)
    out_file.write(out_img)
writer.close()