はじめに
CoreMLを使うとTensorFlow等で学習したモデルをコンバートして、MacやiOSデバイス上で動作させることができる。本記事では以前記事にしたGoogleColab上で学習させた超解像のモデルをMacのアプリ上で利用する手順について記す。
環境
| モジュール | バージョン |
|---|---|
| TensorFlow | 2.4.1 |
| CoreMLTools | 4.1 |
| Xcode | 12.5 |
モデルの学習
これは前述の記事通りなので多くは記さない。
モデルについては若干の違いがあるので、コードを載せておく。
def upsample(x, scale, num_filters):
if 2 <= scale <= 3 :
x = layers.Conv2D(num_filters * (scale ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, scale))(x)
elif scale == 4:
x = layers.Conv2D(num_filters * (2 ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, 2))(x)
x = layers.Conv2D(num_filters * (2 ** 2), 3, padding='same')(x)
x = layers.Lambda(lambda y: tf.nn.depth_to_space(y, 2))(x)
return x
class Normalize(tf.keras.layers.Layer):
def __init__(self, denormalize=False, **kwargs):
super().__init__(**kwargs)
self.denormalize = denormalize
self.DIV2K_RGB_MEAN = np.array([0.4488, 0.4371, 0.4040])* 255
def build(self, input_shape):
super().build(input_shape)
def call(self, x):
rgb_mean=self.DIV2K_RGB_MEAN
if self.denormalize:
x = (x * 127.5) + rgb_mean
x = tf.clip_by_value(x, 0.0, 255.0)
else:
x = (x - rgb_mean) / 127.5
return x
def compute_output_shape(self, input_shape):
return input_shape
def create_model(num_filters=64, num_res_blocks=16, scale=2, shape=(None,None,3), scaling_factor=None, activation='relu'):
inputs = layers.Input( shape , dtype=tf.float32, name='input')
CONV_KWARGS ={
'padding':'same',
'kernel_initializer': tf.keras.initializers.VarianceScaling(scale=0.0001/num_res_blocks,mode='fan_in', distribution='truncated_normal')
}
if activation=='leakyrelu':
activation = tf.keras.layers.LeakyReLU(alpha=0.1)
def res_block(x_in, filters, scaling=None):
x = layers.Conv2D(filters, 3, activation=activation, **CONV_KWARGS)(x_in)
x = layers.Conv2D(filters, 3, **CONV_KWARGS)(x)
if scaling:
x = layers.Lambda(lambda t: t * scaling)(x)
x = layers.Add()([x_in, x])
return x
lowres = Normalize()(inputs)
x = lowres = layers.Conv2D(num_filters, 3, **CONV_KWARGS)(lowres)
for i in range(num_res_blocks):
x = res_block(x, num_filters, scaling=scaling_factor)
x = layers.Conv2D(num_filters, 3, **CONV_KWARGS)(x)
lowres = layers.Add()([x, lowres])
highres = upsample(lowres, scale, num_filters)
highres = layers.Conv2D(3, 3, **CONV_KWARGS)(highres)
output = Normalize(denormalize=True)(highres)
return tf.keras.models.Model(inputs, output)
基本的にcreate_model()をデフォルトのまま呼び出せば良いが、activationはLeakyReLUを使った方が良かったので実験はそれで行った。その他エポック数を伸ばしたり、学習データ生成方法を少し変えたりしたが、詳細は記事の主旨とは関係ないので省略する。
以下、2倍の超解像について生成画像を比較。(左が今回のモデル、右がLANCZOS)
超解像のモデルとしてはEDSRのサブセットのようなものだが、他のモデルとのパラメータ数とPSNRの比較も載せておく。(出典はこちらの論文)
| モデル | パラメータ数 | Set5 | Set14 | BSD100 | Urban100 |
|---|---|---|---|---|---|
| (Bicubic) | - | 33.68 | 30.24 | 29.56 | 26.88 |
| VDSR | 665k | 37.53 | 33.05 | 31.90 | 30.77 |
| EDSR | 43000k | 38.11 | 33.92 | 32.32 | 32.93 |
| 本モデル | 1370k | 38.01 | 33.79 | 32.12 | 31.58 |
コンバート
CoreMLToolsを使ってモデルデータをコンバートする。
CoreMLToolsはGoogleColab上にインストールして実施した。(筆者はM1 Macを所有しているのだが、そこにはうまくインストールできなかった)
下記のようなセルを実行すればコンバートできる。
!pip install coremltools
import coremltools as ct
model = create_model(num_filters=64, num_res_blocks=16, scale=2, shape=(64,64,3), activation='leakyrelu')
model.load_weights(SAVE_DIR + 'F64R16_None_2X_leakyrelu.weights.h5')
mlmodel = ct.convert(model)
coreml_model_path = 'SuperResolution.mlmodel'
mlmodel.save(SAVE_DIR+coreml_model_path)
事前に学習済みの重みを保存してあるものとする。
モデル作成時に入力のshapeを指定する必要がある。この指定は学習時には必要ではないし、コンバート自体も指定しなくても成功するが、アプリに組み込む際にうまく処理できないようだ。ここでは64x64のサイズを指定した。(ここは学習時のパッチサイズと違っていても問題ない)
以前のCoreMLではLambdaやカスタムレイヤーがあると、組み込み時に結構面倒なことになっていたようだが、今回のモデルはLambdaやカスタムレイヤーを含むにもかかわらず、特に問題なくコンバートできた。
アプリ側組み込み
上記で作成した'SuperResolution.mlmodel'をXcodeのプロジェクトに追加して、以下のようなクラスでMacのアプリに組み込んだ。
import Foundation
import CoreML
import Cocoa
class SuperResolver{
private func cgImageFromRGB(rgb: UnsafeMutablePointer<UInt8>, width: Int, height: Int, bytesPerPixel:Int ) -> CGImage? {
let bitsPerComponent = 8
let bitsPerPixel = bytesPerPixel*bitsPerComponent
let count = height * width * bytesPerPixel
let releasePixelData: CGDataProviderReleaseDataCallback = { (info: UnsafeMutableRawPointer?, data: UnsafeRawPointer, size: Int) -> () in
print(size)
return
}
let provider = CGDataProvider(dataInfo: nil,data: rgb, size: count, releaseData: releasePixelData)
return CGImage(
width: width,
height: height,
bitsPerComponent: bitsPerComponent,
bitsPerPixel:bitsPerPixel,
bytesPerRow: width * bytesPerPixel,
space: CGColorSpaceCreateDeviceRGB(),
bitmapInfo: [],
provider: provider!,
decode: nil,
shouldInterpolate: false,
intent: CGColorRenderingIntent.defaultIntent)
}
private func clamp<T: Comparable>(_ value: T, _ minValue: T, _ maxValue: T) -> T {
if value < minValue {
return minValue
}
if value > maxValue {
return maxValue
}
return value
}
private func getPatch(cgImage:CGImage, origin:CGPoint, size:CGSize) -> [Float]
{
let bytesPerRow = cgImage.bytesPerRow
let bytesPerPixel = cgImage.bitsPerPixel/8
let pixelData = cgImage.dataProvider!.data! as Data
var rgb_buf = Array<Float>(repeating:Float(0), count:Int(size.height)*Int(size.width)*3)
var dstOffset : Int = 0
for y1 in 0..<Int(size.height) {
let y = clamp(y1 + Int(origin.y), 0, cgImage.height-1)
for x1 in 0..<Int(size.width){
let x = clamp(x1 + Int(origin.x), 0, cgImage.width-1)
let pixelOffset = (bytesPerRow * y) + ( x * bytesPerPixel)
rgb_buf[dstOffset] = Float(pixelData[pixelOffset])
rgb_buf[dstOffset+1] = Float(pixelData[pixelOffset+1])
rgb_buf[dstOffset+2] = Float(pixelData[pixelOffset+2])
dstOffset += 3
}
}
return rgb_buf
}
private func updateImageArray( dst:UnsafeMutablePointer<UInt8>, dstWidth:Int, dstHeight:Int, patch:MLMultiArray, origin:CGPoint, size:CGSize){
for y1 in 0..<Int(size.height) {
let dstY = Int(origin.y) + y1
if dstY >= dstHeight {
break
}
if dstY<0 {
continue
}
for x1 in 0..<Int(size.width) {
let dstX = Int(origin.x) + x1
if dstX >= dstWidth {
break
}
if dstX<0 {
continue
}
let dstOffset = (dstY*dstWidth + dstX)*3
let srcOffset = (y1*Int(size.width) + x1)*3
dst[dstOffset ] = UInt8(truncating: patch[srcOffset])
dst[dstOffset+1] = UInt8(truncating: patch[srcOffset+1])
dst[dstOffset+2] = UInt8(truncating: patch[srcOffset+2])
}
}
}
func resolve( img: NSImage ) -> NSImage{
let cgImage = img.cgImage(forProposedRect: nil, context: nil, hints: nil)!
let imgWidth = Int(img.size.width)
let imgHeight = Int(img.size.height)
let config = MLModelConfiguration()
config.computeUnits = MLComputeUnits.all
let model = try! SuperResolution(configuration:config)
let scale = 2
let count = imgWidth * imgHeight * scale * scale * 3
let srImageArray = UnsafeMutablePointer<UInt8>.allocate(capacity:count)
let patchSize = 64
let mlArray = try! MLMultiArray(shape: [1,NSNumber(value:patchSize),NSNumber(value:patchSize),3], dataType: MLMultiArrayDataType.float32)
for y in stride(from:0, to:imgHeight, by:patchSize){
for x in stride(from:0, to:imgWidth, by:patchSize){
let patchPixel = getPatch(cgImage: cgImage, origin:CGPoint(x:x,y:y), size:CGSize(width:patchSize,height:patchSize))
mlArray.dataPointer.initializeMemory(as: Float.self, from:patchPixel, count:patchPixel.count)
let output = try! model.prediction(input: SuperResolutionInput(input:mlArray))
updateImageArray(dst:srImageArray, dstWidth:imgWidth*scale, dstHeight:imgHeight*scale,
patch:output.Identity, origin:CGPoint(x:x*scale,y:y*scale), size:CGSize(width:patchSize*scale,height:patchSize*scale) )
}
}
let cgImageSR = cgImageFromRGB(rgb: srImageArray, width: imgWidth*scale, height: imgHeight*scale, bytesPerPixel: 3)!
return NSImage(cgImage: cgImageSR, size: CGSize(width: cgImageSR.width, height: cgImageSR.height))
}
MacのアプリなのでNSImageを使っているが、iOSの場合はUIImageを使うことになるだろう。
このクラスのresolveを呼び出して高解像度画像を得ることになる。
処理としては、以下のようなことをしている。
- SuperResolutionインスタンス生成
- 全体画像用のバッファ生成
- パッチの切り出し
- MLMultiArrayにデータをセット
- SuperResolutionのpredictionを呼び出して高解像度画像取得
- 全体画像のバッファを更新
- 全領域終了するまでパッチ切り出しからの処理を繰り返す
- 全体画像用バッファから画像を生成
モデルの入力サイズが決まっているので、それに合わせて64x64で領域を切り出して処理している。
モデルの入力はMLMultiArrayなので画像データから変換して入力している。コンバートの設定次第ではImageを直接入力できるようだが、今回はデフォルトのままなのでFloatに変換して入力した。pytorchのモデルをコンバートした場合は、色のチャンネルが先に来るので、この辺のコードは修正する必要があると思われる。
単純に64x64に区切って処理をしているので、境界で不連続になるのではないか、と思っていたのだが、実際にやってみると視覚的に気になるようなところは特に出なかったので、簡単な実装のままでよしとした。厳密には境界付近はオーバーラップさせた方がPSNR的には良い結果が得られると思われる。また、回転や反転等した結果も混合させるとセルフアンサンブルとなって向上するだろう。ただし、処理時間は当然増える。
ここで掲載した実装では毎回SuperResolutionのインスタンスを生成しているのだが、これはメモリの占有サイズを減らすためで、連続的に処理するなら一度生成した後で保持して使用し続ける方が処理時間的には良い。
今回の実装をM1 Mac Miniで実行すると、250x240の画像を２倍する処理で約0.25秒だった。モデルのインスタンスを保持して使いまわした場合は、初回を除くと約0.2秒となった。
まとめ
CoreMLTootsを使って、TensorFlowのモデルをコンバートしてMac/iOS上で動作させるアプリの作成手順を示した。
コンバーターは以前から改良されているようで、Lambdaやカスタムレイヤーを使用していても、特に問題なかった。ただし、もっと複雑なモデルでは問題になる可能性もあるので、注意は必要と思われる。