Video Style Transfer

from neural_style_transfer import style_transfer
from tqdm.notebook import tqdm
from PIL import Image
import os
import cv2
%load_ext autoreload
%autoreload 2

Processamento do vídeo

Carregamos o vídeo como um tensor de imagens utilizando dos métodos do openCV. Iremos iterar por esse tensor e realizar o treinamento para cada imagem individualmente.

output_path = '/content/drive/MyDrive/video_style_transfer/output_cat/'
video_path = '/content/drive/MyDrive/video_style_transfer/cat_video.mp4'
style_path = '/content/drive/MyDrive/video_style_transfer/stary_night.jpg'

vidcap = cv2.VideoCapture(video_path)
success,image = vidcap.read()
count = 0
while success:
    image = Image.fromarray(image)
    model = style_transfer(image, style = style_path, 
                           output_path = output_path + 'image_'+ str(count)+'.jpg',
                           iterations = 10, cuda = True, verbose = False)
    _ = model.train()
    success,image = vidcap.read()
    count += 1

Criação do vídeo

Agora, com a pasta de imagens, vamos utilizar da bibilioteca openCV2 para criar um arquivo mp4.

def create_video_from_images(folder, out):
    files = os.listdir(folder)
    files.sort(key = lambda x : int(x[x.find('_')+1:x.find('.')]))
    size = (224, 224)
    out = cv2.VideoWriter(out+'.mp4',cv2.VideoWriter_fourcc(*'MP4V'), 24, size)
    for f in tqdm(files):
        im = cv2.imread(folder+'/'+f)
        im = cv2.rotate(im, cv2.cv2.ROTATE_90_CLOCKWISE) 
        out.write(im)

create_video_from_images(output_path, '/content/drive/MyDrive/video_style_transfer/cat_style')