🎨 Используем CycleGAN для применения стиля к видео, полученному с веб-камеры

bash ./scripts/download_cyclegan_model.sh style_monet_pretrained
bash ./scripts/download_cyclegan_model.sh style_ukiyoe_pretrained
bash ./scripts/download_cyclegan_model.sh style_cezanne_pretrained
bash ./scripts/download_cyclegan_model.sh style_vangogh_pretrained

import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
import cv2
import torch
import numpy as np

if __name__ == '__main__':
    opt = TestOptions().parse()  # получаем опции теста
    # прописываем значения по умолчанию для теста
    opt.num_threads = 0   # тестовый код поддерживает только num_threads = 0
    opt.batch_size = 1    # тестовый код поддерживает только batch_size = 1
    opt.serial_batches = True  # запретить перетасовку данных; закомментируйте эту строку, если нужны результаты для случайно выбранных изображений 
    opt.no_flip = True    # нет поворотов; закомментируйте эту строку, если нужно поворачивать изображения
    opt.display_id = -1   # тестовый код сохраняет результаты в файл HTML
    model = create_model(opt)      # создать модель по заданному opt.model и другим параметрам
    model.setup(opt)               # обычные настройки: конвейеры загрузки и печати; создание scheduler'ов
    if opt.eval:
        model.eval()

# начинаем получение видео / настройка веб-камеры
    webcam = cv2.VideoCapture(0)
    # Проверим, удачно ли открылась камера
    if not webcam.isOpened():
        raise IOError("Cannot open webcam")

 # CycleGan принимает данные в виде словаря
    # проще выполнить это требование, чем переделывать
    # начинаем бесконечный цикл - читаем фреймы с веб-камеры, пока пользователь не прервет цикл с клавиатуры
    data = {"A": None, "A_paths": None}
    while True:

        #ret - это bool, который вернула cap.read() -> независимо от того, удалось ли считать фрейм
        #если считали удачно, сохраняем его в frame
        ret, frame = webcam.read()

        #изменяем размеры frame
        frame = cv2.resize(frame, (256,256), interpolation=cv2.INTER_AREA)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        #модель рассчитывает получить batchsize * channels * h * w
        #добавляем измерение для размера пакета 
        frame = np.array([frame])
        #теперь форма - batchsize * channels * h * w
        frame = frame.transpose([0,3,1,2])

        #переводим массив numpy в тензор
        #данные должны быть тензором для совместимости с запущенной моделью, ожидающей floatTensor'ы
        data['A'] = torch.FloatTensor(frame)
        
        model.set_input(data)  # распаковать данные из загрузчика 
        model.test()

        #получаем только сгенерированное изображение - ищем в словаре ключ "fake"
        result_image = model.get_current_visuals()['fake']
        #используем tensor2im, предоставленный в файле util
        result_image = util.tensor2im(result_image)
        result_image = cv2.cvtColor(np.array(result_image), cv2.COLOR_BGR2RGB)  
        result_image = cv2.resize(result_image, (512, 512))  

    cv2.imshow('style', result_image)

    #код ASCII клавиши Esc - 27.
    c = cv2.waitKey(1)
    if c == 27:
        break
  cap.release()
  cv2.destroyAllWindows()

python webcam.py  --name MODEL_NAME --model test --preprocess none --no_dropout

#настройка циклического переключения стилей
style_models = ['style_monet_pretrained', 'style_vangogh_pretrained', 'style_ukiyoe_pretrained', 'style_cezanne_pretrained']
style_model_index = 0

#настройка текста
# шрифт
font = cv2.FONT_HERSHEY_SIMPLEX 
# расположение
org = (0, 25) 
# масштаб шрифта
fontScale = 1
# Синий цвет фона 
color = (255, 255, 255) 
# Толщина линий 2 пикселя
thickness = 2

"""Тестовый скрипт общего назначения для перевода изображений.

После обучения вашей модели с помощью train.py, вы можете использовать этот скрипт для ее тестирования.
Он загрузит сохраненную модель из '--checkpoints_dir' и сохранит результаты в '--results_dir'.

Сначала он создает модель и набор данных в соответствии с опциями. Некоторые параметры жестко прописаны в коде.
Затем он запускает перевод для '--num_test' изображений и сохраняет результаты в файл HTML.

Пример (Сначала надо обучить модели или загрузить предобученные с нашего сайта):
    Тестировать модель CycleGAN (в обе стороны):
        python test.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan

    Тестировать модель CycleGAN (только в одну сторону):
        python test.py --dataroot datasets/horse2zebra/testA --name horse2zebra_pretrained --model test --no_dropout

    Опция '--model test' используется для генерации результатов CycleGAN только в одну сторону.
    Эта опция автоматически устанавливает '--dataset_mode single', при которой изображения загружаются только из одного набора.
    Напротив, использование '--model cycle_gan' требует загрузки и генерации результатов в обе стороны,
    что иногда не нужно. Результаты будут сохранены в ./results/.
    Use '--results_dir <directory_path_to_save_result>' to specify the results directory.

    Test a pix2pix model:
        python test.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA

См. больше опций в options/base_options.py и options/test_options.py.
См. подсказки для обучения и тестов на: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md
Часто задаваемые вопросы см: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md
"""
import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
import cv2
import torch
import numpy as np


if __name__ == '__main__':
    opt = TestOptions().parse()  # загружаем тестовые опции
    # жестко пропишем некоторые параметры для теста
    opt.num_threads = 0   # тестовый код поддерживает только num_threads = 0
    opt.batch_size = 1    # тестовые код поддерживает только batch_size = 1
    opt.serial_batches = True  # запретить перетасовку данных; закомментируйте эту строку, если нужен перевод случайно выбранных изображений.
    opt.no_flip = True    # нет поворотов; закомментируйте эту строку, если нужен результат на повернутых изображениях.
    opt.display_id = -1   # тестовый код сохраняет результаты в файл HTML.
    model = create_model(opt)      # создать модель по заданному opt.model и другим параметрам
    model.setup(opt)               # обычные настройки: конвейеры загрузки и печати; создание scheduler'ов
    if opt.eval:
        model.eval()
    
    #начать съемку видео/настроить камеру
    webcam = cv2.VideoCapture(0)
    # Проверить, корректно ли открылась камера
    if not webcam.isOpened():
        raise IOError("Cannot open webcam")
    
    #циклический проход по стилям
    style_models = ['style_monet_pretrained', 'style_vangogh_pretrained', 'style_ukiyoe_pretrained', 'style_cezanne_pretrained']
    style_model_index = 0

    #настройка текста
    # шрифт
    font = cv2.FONT_HERSHEY_SIMPLEX 
    # начальная точка
    org = (0, 25) 
    # масштаб шрифта
    fontScale = 1
    # Синий цвет фона
    color = (255, 255, 255) 
    # Толщина линий 2 пикселя
    thickness = 2


    #CycleGan принимает данные в виде словаря
    # проще выполнить это требование, чем переделывать
    # начинаем бесконечный цикл - читаем фреймы с веб-камеры, пока пользователь не прервет цикл с клавиатуры
    data = {"A": None, "A_paths": None}
    while True:

        #ret - это bool, который вернула cap.read() -> независимо от того, удалось ли считать фрейм
        #если считали удачно, сохраняем его в frame
        ret, frame = webcam.read()

        #изменяем размеры frame
        frame = cv2.resize(frame, (256,256), interpolation=cv2.INTER_AREA)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        #модель рассчитывает получить batchsize * channels * h * w
        #добавляем измерение для размера пакета
        frame = np.array([frame])
        #теперь форма - batchsize * channels * h * w
        frame = frame.transpose([0,3,1,2])

        #переводим массив numpy в тензор
        #данные должны быть тензором для совместимости с запущенной моделью, ожидающей floatTensor'ы
        data['A'] = torch.FloatTensor(frame)
        
        model.set_input(data)  # распаковать данные из загрузчика
        model.test()

        #получаем только сгенерированное изображение - ищем в словаре ключ "fake"
        result_image = model.get_current_visuals()['fake']
        #используем tensor2im, предоставленный в файле util
        result_image = util.tensor2im(result_image)
        result_image = cv2.cvtColor(np.array(result_image), cv2.COLOR_BGR2RGB)  
        result_image = cv2.resize(result_image, (512, 512))      
        result_image = cv2.putText(result_image, str(opt.name)[6:-11], org, font,  
                   fontScale, color, thickness, cv2.LINE_AA)   
        cv2.imshow('style', result_image)

        #ASCII value of Esc is 27.
        c = cv2.waitKey(1)
        if c == 27:
            break
        if c == 99:
            if style_model_index == len(style_models):
                style_model_index = 0
            opt.name = style_models[style_model_index]
            style_model_index += 1
            model = create_model(opt)      # создаем модель с заданными опциями create a model given opt.model and other options
            model.setup(opt) 
      
        
    cap.release()
    cv2.destroyAllWindows()
        

python webcam.py  --name MODEL_NAME --model test --preprocess none --no_dropout