ケーススタディGFPGAN
GFPGANは品質の良くない顔画像を綺麗にできる。
このPytorchモデルをCoreMLで画像出力できるようにしていく。
そのままではうまく画像にできない
CoreMLToolsを使ってPytorchモデルをCoreMLモデルに変換するとき、以下のようにして出力を画像に設定することができるが、通常このままではうまく画像が出力されない。
import coremltools as ct
dummy_input = torch.randn((1,3,512,512)).cuda()
jit_model = torch.jit.trace(model,dummy_input)
coreml_model = ct.convert(
jit_model,
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT32,
compute_units=ct.ComputeUnit.CPU_AND_GPU,
inputs=[
ct.ImageType(name="image",
shape=dummy_input.shape,
bias=[-1,-1,-1],
scale=1/127.5)
],
outputs=[ct.ImageType(name="output")])
真っ黒が出力された結果↓ かなしい。
なぜかというと、
通常Pytorchモデル実行時には、モデル推論後に結果の後処理を行なって画像にしており、そのままモデルだけをCoreMLモデル変換すると、後処理を含まないからである。
後処理が必要
たとえば、torchの出力は-1~1に正規化されていたりするが、画像のピクセルレンジは0~255なので、そのまま画像にすると値が小さすぎて真っ黒な画像になる。
なので、結果の値に127.5をかけて、127.5を足すことで0~255に戻している。
# pytorchの後処理
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_((-1,1))
_tensor = (_tensor - (-1)) / (1 - (-1))
img_np = _tensor.numpy()
img_np = img_np.transpose(1, 2, 0)
img_np = (img_np * 255.0).round()
これはBasicSRのtensor2imgを使った後処理で、わかりにくいですが、要するに-1~1を0~1にして２５５をかけているということだと思います。
以下と同等だと思う。
output = torch.clamp(_tensor * 127.5 + 127.5, min=0, max=255)
CV2では（512,512,3)として扱う必要がありますが、CoreMLモデルに変換する際は4次元のshape（1,3,512,512）でいいので、ここではsqueezeとtransposeは必要ありません。
Pythonでラップモデルを作ってから変換
で、どうやって後処理をiOSで使えばいいかというと、swiftでmultiarrayを扱ってもいいのですが、CoreMLモデルに含めてしまいましょう。
変換前に後処理を追加したモデルクラスを作って変換します。
class CoreMGFPGAN(torch.nn.Module):
def __init__(self, gfpgan):
super(CoreMGFPGAN, self).__init__()
self.gfpgan = gfpgan
def forward(self, image):
gfpgan_out = self.gfpgan(image)
output = torch.clamp(gfpgan_out * 127.5 + 127.5, min=0, max=255)
return output
model = CoreMGFPGAN(gfpgan).eval()
あとは普通に変換すればいいだけ。
from coremltools.converters.mil.input_types import ColorLayout
ex = torch.randn((1,3,512,512)).cuda()
jit_model = torch.jit.trace(model,ex)
import coremltools as ct
coreml_model = ct.convert(
jit_model,
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT32,
compute_units=ct.ComputeUnit.CPU_AND_GPU,
inputs=[
ct.ImageType(name="image",
shape=ex.shape,
bias=[-1,-1,-1],
scale=1/127.5)
],
outputs=[ct.ImageType(name="output")])
これで画像が出力されます。
変換例はこちら。
🐣
フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。
もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…
いずれも中間コストを省いたリーズナブルな価格でお請けできます。
お仕事のご相談はこちらまで
rockyshikoku@gmail.com
機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。