近日,已使用多年的人教版小學數學教材中的插畫引發社會各界人士爭議。咱程序員也沒有手繪插畫能力,但咱可以借助強大的深度學習模型將視頻轉動漫。本文目標是讓任何具有python
語言基本能力的程序員,實現短視頻轉動漫效果。示例效果如下:
整個實現流程如下:
- 讀取視頻幀
- 將每一幀圖像轉為動漫幀
- 將轉換後的動漫幀轉為視頻
難點在於如何將圖像轉為動漫效果。這裡我們使用基於深度學習的動漫效果轉換模型,考慮到許多讀者對這塊不了解,因此我這邊准備好了源碼和模型,直接調用即可。不想看文章細節的可以直接拖到文章末尾,獲取源碼。
為了讓讀者不關心深度學習模型,已經為讀者准備好了轉換後的onnx
類型模型。接下來按順序介紹運行onnx
模型流程。
onnxruntime
庫pip install onnxruntime
如果想要用GPU
加速,可以安裝GPU
版本的onnxruntime
:
pip install onnxruntime-gpu
需要注意的是:
onnxruntime-gpu
的版本跟CUDA
有關聯,具體對應關系如下:
當然了,如果用CPU
運行,那就不需要考慮那麼多啦。考慮到通用性,本文全部以CPU
版本onnxruntime
。
先導入onnxruntime
庫,創建InferenceSession
對象,調用run函數。如下所示
import onnxruntime as rt
sess = rt.InferenceSession(MODEL_PATH)
inp_name = sess.get_inputs()[0].name
out = sess.run(None, {
inp_name: inp_image})
具體到我們這裡的動漫效果,實現細節如下:
import cv2
import numpy as np
import onnxruntime as rt
# MODEL = "models/anime_1.onnx"
MODEL = "models/anime_2.onnx"
sess = rt.InferenceSession(MODEL)
inp_name = sess.get_inputs()[0].name
def infer(rgb):
rgb = np.expand_dims(rgb, 0)
rgb = rgb * 2.0 / 255.0 - 1
rgb = rgb.astype(np.float32)
out = sess.run(None, {
inp_name: rgb})
out = out[0][0]
out = (out+1)/2*255
out = np.clip(out, 0, 255).astype(np.uint8)
return out
def preprocess(rgb):
pad_w = 0
pad_h = 0
h,w,__ = rgb.shape
N = 2**3
if h%N!=0:
pad_h=(h//N+1)*N-h
if w%2!=0:
pad_w=(w//N+1)*N-w
# print(pad_w, pad_h, w, h)
rgb = np.pad(rgb, ((0,pad_h),(0, pad_w),(0,0)), "reflect")
return rgb, pad_w, pad_h
其中, preprocess
函數確保輸入圖像的寬高是8
的整數倍。這裡主要是因為考慮到深度學習模型有下采樣,確保每次下采樣能被2
整除。
這裡使用Opencv
庫,提取視頻中每一幀並調用回調函數將視頻幀回傳。在將圖片轉視頻過程中,通過定義VideoWriter
類型變量WRITE
確保唯一性。具體實現代碼如下:
import cv2
from tqdm import tqdm
WRITER = None
def write_frame(frame, out_path, fps=30):
global WRITER
if WRITER is None:
size = frame.shape[0:2][::-1]
WRITER = cv2.VideoWriter(
out_path,
cv2.VideoWriter_fourcc(*'mp4v'), # 編碼器
fps,
size)
WRITER.write(frame)
def extract_frames(video_path, callback):
video = cv2.VideoCapture(video_path)
num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
for _ in tqdm(range(num_frames)):
_, frame = video.read()
if frame is not None:
callback(frame)
else:
break
Python學習實戰
動漫
,獲取完整源碼。如果您覺得本文有幫助,辛苦您點個不需花錢的贊,您的舉手之勞將對我提供了無限的寫作動力! 也歡迎關注我的公眾號:Python學習實戰, 第一時間獲取最新文章。