Разработка
Как сделать аналог Prisma на Fast Style Transfer, CoreML и TensorFlow
Разработчик Майкл Рамос рассказал, как можно сделать аналог Prisma своими руками.
Вступление
Основная часть этого туториала взята из блога Prisma Lab и их работы с PyTorch. Однако мы будем использовать TensorFlow для создания моделей, в частности Fast Style Transfer от Логана Энгстрема. Результатом этого урока будет iOS-приложение, которое будет запускать модели TensorFlow при помощи CoreML. Вот репозиторий в GitHub, в котором содержатся все нужные дополнения.
Благодаря чему это стало возможным:
- Стэнфордское исследование
- Fast Style Transfer
- CoreML от Apple (не поддерживает TensorFlow)
- Google выпускает TensorFlow Lite (не поддерживает CoreML)
- Google выпускает поддержку CoreML (неполная поддержка)
- Мы внесли несколько дополнений и создали готовое решение
Инструменты
Модели: мы будем использовать уже обученные модели FST, но так же хорошо сработают и кастомные модели, вам только нужно будет сделать несколько небольших изменений, которые я упомяну. Скачать модели.
TensorFlow: При использовании FST лучше всего использовать TensorFlow 1.0.0, а при использовании TensorFlow-CoreML — версию 1.1.0 или новее (для этого урока GPU не понадобится).
iOS: 11
Xcode: 9
Python: 2.7
Подготовка
Нам нужно сделать несколько подготовительных шагов, так как FST является в большей степени решением для исследований, а не для повторного использования и создания продукта — в нем нет соглашения об именовании или изображения в качестве результата.
Шаг 1: Нам нужно придумать название выходного изображения, иначе TensorFlow сгенерирует его автоматически. Мы можем сделать это в скрипте evaluate.py.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# function ffwd,line 93 | |
# https://github.com/lengstrom/fast-style-transfer/blob/master/evaluate.py#L93 | |
preds = transform.net(img_placeholder) | |
# !! Add This !! | |
print(preds) |
После этого мы можем запустить скрипт, чтобы увидеть результат. Я здесь использую уже обученную модель wave (если вы используете свои модели, параметр checkpoint просто должен быть директорией, где хранятся ваши мета-файлы и входные данные).
$ python evaluate.py --checkpoint wave.ckpt --in-path inputs/ --out-path outputs/
> Tensor(“add_37:0”, shape=(20, 720, 884, 3), dtype=float32, device=/device:GPU:0)
Здесь имеет значение только название выходного узла, то есть add_37. Это имеет смысл, так как последний неназванный оператор в сети — это сложение.
Шаг 2: Нам нужно внести ещё несколько изменений в evaluate.py, чтобы сохранить изображение на диск. Обратите внимание, что если вы используете ваши собственные модели, то вам нужно будет добавить код работы с каталогом контрольных точек по сравнению с одним файлом контрольной точки.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# function ffwd, line 98 | |
# https://github.com/lengstrom/fast-style-transfer/blob/master/evaluate.py#L98 | |
if os.path.isdir(checkpoint_dir): | |
ckpt = tf.train.get_checkpoint_state(checkpoint_dir) | |
if ckpt and ckpt.model_checkpoint_path: | |
saver.restore(sess, ckpt.model_checkpoint_path) | |
########## add this for pre-trained models ########### | |
frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess,sess.graph_def,['add_37']) | |
with open('output_graph.pb', 'wb') as f: | |
f.write(frozen_graph_def.SerializeToString()) | |
##################################################### | |
else: | |
raise Exception("No checkpoint found…") | |
else: | |
saver.restore(sess, checkpoint_dir) | |
########## add this for custom models ########### | |
frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess,sess.graph_def,['add_37']) | |
with open('output_graph.pb', 'wb') as f: | |
f.write(frozen_graph_def.SerializeToString()) | |
##################################################### |
Шаг 3: Теперь мы запустим evaluate.py для модели и увидим, что ваш файл с изображением сохранен. При обучении моделей вы, вероятно, использовали больше одного образца для тренировки (batch size), а также GPU, однако CoreML принимает только граф с размером 1 и оптимизацией для CPU. Обратите внимание на команду оценки для настройки.
$ python evaluate.py --checkpoint wave/wave.ckpt --in-path inputs/ --out-path outputs/ --device “/cpu:0” --batch-size 1
Отлично, мы создали output_graph.pb и можем приступать к конверсии в CoreML.
Конверсия в CoreML
Благодаря Google теперь существует конвертер из TensorFlow в CoreML. Это отлично, но решение новое, и в нем не хватает некоторых важных операций TensorFlow, например, power.
Шаг 1: Наша модель не будет конвертироваться без поддержки power, но, к счастью, инструменты CoreML предоставляют унарную конверсию, которая поддерживает power. Нам нужно добавить этот код в решение на TensorFlow.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# tfcoreml src | |
# file1 : _interpret_shapes.py | |
# | |
# in the _SHAPE_TRANSLATOR_REGISTRY we need to add the Pow operation | |
_SHAPE_TRANSLATOR_REGISTRY = { | |
… previous keys … | |
# add this: | |
'Pow': _identity, | |
} | |
# file 2: _ops_to_layers.py | |
# | |
# in the _OP_REGISTRY to add the Pow operation | |
_OP_REGISTRY = { | |
… previous keys … | |
# add this: | |
'Pow': _layers.pow | |
} | |
# file 3: _layers.py | |
# | |
# in the _layers we need to define the conversion | |
def pow(op, context): | |
const_name = compat.as_bytes(op.inputs[1].name) | |
const_val = context.consts[const_name] ## Note: this is .5 here, you can toy around if you want | |
input_name = compat.as_bytes(op.inputs[0].name) | |
output_name = compat.as_bytes(op.outputs[0].name) | |
context.builder.add_unary(output_name, input_name, output_name, 'power', alpha=const_val) | |
context.translated[output_name] = True |
Шаг 2: Создайте и запустите скрипт конверсии.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import tfcoreml as tf_converter | |
tf_converter.convert(tf_model_path = 'output_graph.pb', | |
mlmodel_path = 'model_name.mlmodel', | |
output_feature_names = ['add_37:0'], | |
## Note found this after running a conversion the first time | |
image_input_names = ['img_placeholder__0']) |
$ python convert.py
Актуальный конвертер CoreML не предоставляет возможности вывода изображений из модели. Изображения представлены массивами NumPy (многомерные массивы), которые являются фактическим результатом и скомпилированы для нестандартного типа MultiArray в Swift. Я поискал в интернете и смог получить код, который оценивал выходные данные от CoreML, а затем преобразовывал их в изображения.
Шаг 3: создайте и запустите скрипт трансформации выходных данных в модели my_model.mlmodel.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import coremltools | |
def convert_multiarray_output_to_image(spec, feature_name, is_bgr=False): | |
""" | |
Convert an output multiarray to be represented as an image | |
This will modify the Model_pb spec passed in. | |
Example: | |
model = coremltools.models.MLModel('MyNeuralNetwork.mlmodel') | |
spec = model.get_spec() | |
convert_multiarray_output_to_image(spec,'imageOutput',is_bgr=False) | |
newModel = coremltools.models.MLModel(spec) | |
newModel.save('MyNeuralNetworkWithImageOutput.mlmodel') | |
Parameters | |
———- | |
spec: Model_pb | |
The specification containing the output feature to convert | |
feature_name: str | |
The name of the multiarray output feature you want to convert | |
is_bgr: boolean | |
If multiarray has 3 channels, set to True for RGB pixel order or false for BGR | |
""" | |
for output in spec.description.output: | |
if output.name != feature_name: | |
continue | |
if output.type.WhichOneof('Type') != 'multiArrayType': | |
raise ValueError("%s is not a multiarray type" % output.name) | |
array_shape = tuple(output.type.multiArrayType.shape) | |
channels, height, width = array_shape | |
from coremltools.proto import FeatureTypes_pb2 as ft | |
if channels == 1: | |
output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('GRAYSCALE') | |
elif channels == 3: | |
if is_bgr: | |
output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('BGR') | |
else: | |
output.type.imageType.colorSpace = ft.ImageFeatureType.ColorSpace.Value('RGB') | |
else: | |
raise ValueError("Channel Value %d not supported for image inputs" % channels) | |
output.type.imageType.width = width | |
output.type.imageType.height = height | |
model = coremltools.models.MLModel('my_model.mlmodel') | |
spec = model.get_spec() | |
convert_multiarray_output_to_image(spec,'add_37__0',is_bgr=False) | |
newModel = coremltools.models.MLModel(spec) | |
newModel.save('wave.mlmodel') |
$ python output.py
И… 🎉💥🤙… у нас есть рабочая модель CoreML. Я изменил её название на wave в выходном скрипте, а также размер изображения, чтобы тот соответствовал входным данным.
Приложение iOS
Это не совсем урок, поэтому вы можете поработать с моим репозиторием. Я затрону здесь только самое важное.
Шаг 1: Импортируйте модели в свой проект Xcode.
Шаг 2: После импортирования вы сможете инициализировать свои модели:
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
private let models = [ | |
wave().model, | |
udnie().model, | |
rain_princess().model, | |
la_muse().model | |
] |
Шаг 3: Создайте класс для входного параметра модели, MLFeatureProvider. img_placeholder — это входные данные, которые определены в скрипте оценки.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
private let models = [ | |
wave().model, | |
udnie().model, | |
rain_princess().model, | |
la_muse().model | |
] |
Шаг 4: Теперь мы можем вызвать модель в нашем коде.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
private func stylizeImage(cgImage: CGImage, model: MLModel) -> CGImage { | |
// size can change here if you want, remember to run right sizes in the fst evaluating script | |
let input = StyleTransferInput(input: pixelBuffer(cgImage: cgImage, width: 883, height: 720)) | |
// model.prediction will run the style model on input image | |
let outFeatures = try! model.prediction(from: input) | |
// we get the image buffer after | |
let output = outFeatures.featureValue(for: "add_37__0")!.imageBufferValue! | |
// remaining code to convert image buffer here ….. | |
} |
Финал: Остальная часть приложения — это настройки и обработка изображений. Ничего нового или связанного с CoreML, поэтому мы не будем это рассматривать. На этой точке у вас будет понимание того, как все работает вместе, и вы сможете вносить свои улучшения. Я думаю, что как минимум можно улучшить граф FST. Мы должны убрать слишком сложные операции, чтобы финальное приложение работало ещё быстрее. Но и сейчас все работает довольно неплохо.
-
Рекламные сети1 месяц назад
Telega.in (Телега. ин) – обзор и отзывы. Маркировка рекламных материалов.
-
Разработка4 недели назад
Осваиваем ViewModel в Android: «можно» и «нельзя» — Часть 2
-
Маркетинг и монетизация1 месяц назад
Мобильная разработка за неделю #491: быстрый старт в растрату карьеры
-
Статьи1 месяц назад
10 упражнений от боли в спине, которые помогут программистам