Разработка
Как сделать аналог 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.
| # 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, чтобы сохранить изображение на диск. Обратите внимание, что если вы используете ваши собственные модели, то вам нужно будет добавить код работы с каталогом контрольных точек по сравнению с одним файлом контрольной точки.
| # 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.
| # 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: Создайте и запустите скрипт конверсии.
| 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.
| 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: После импортирования вы сможете инициализировать свои модели:
| private let models = [ | |
| wave().model, | |
| udnie().model, | |
| rain_princess().model, | |
| la_muse().model | |
| ] |
Шаг 3: Создайте класс для входного параметра модели, MLFeatureProvider. img_placeholder — это входные данные, которые определены в скрипте оценки.
| private let models = [ | |
| wave().model, | |
| udnie().model, | |
| rain_princess().model, | |
| la_muse().model | |
| ] |
Шаг 4: Теперь мы можем вызвать модель в нашем коде.
| 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. Мы должны убрать слишком сложные операции, чтобы финальное приложение работало ещё быстрее. Но и сейчас все работает довольно неплохо.
