RMBG1.4をCore ML モデルに変換する

# ハイクオリティな背景除去をiOSで

![pexels-oleksandr-p-1004359.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/6eeebac8-ae5a-0cfc-75e8-7392795040b6.jpeg)
![example_image_no_bg.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/662c1fa2-4d1a-e4b9-2f76-b32c59871c4b.png)


### オリジナルモデルとCoreMLToolsのインストール

```shell
git clone https://huggingface.co/briaai/RMBG-1.4
cd RMBG-1.4/
pip install -r requirements.txt
pip install coremltools
```

### 方針

正規化した入力を受け取り、グレイスケールのマスク画像を出力するようにします。
Pythonのpostprocessを含めたラップモデルを作り、
preprocessを含めたCore ML変換をする。

PythonのpreprocessとpostProcessは以下。

```python
def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
    if len(im.shape) < 3:
        im = im[:, :, np.newaxis]
    # orig_im_size=im.shape[0:2]
    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear').type(torch.uint8)
    image = torch.divide(im_tensor,255.0)
    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
    return image

def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:
    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)
    ma = torch.max(result)
    mi = torch.min(result)
    result = (result-mi)/(ma-mi)
    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
    im_array = np.squeeze(im_array)
    return im_array
```

preprocessは入力を-0.5~0.5に正規化している。
postprocessは出力を0~１にした後、0~255にしている。
permute操作（配列の次元入れ替え）はcv2で操作するためのものなので、今回は入れない。

オリジナルモデルとpostprocessを含んだラップモデルを作る

```python
from briarmbg import BriaRMBG

net = BriaRMBG()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net = BriaRMBG.from_pretrained("briaai/RMBG-1.4")
net.to(device)

# wrap class including the postprocess
class CoreMLRMBG(torch.nn.Module):
    def __init__(self, net):
        super(CoreMLRMBG, self).__init__()
        self.net = net

    def forward(self, image):
        result = self.net(image)[0][0]
        ma = torch.max(result)
        mi = torch.min(result)
        result = (result-mi)/(ma-mi)
        im_array = (result*255)
        return im_array

model = CoreMLRMBG(net)
````

オリジナルのモデルを丸々入れて、フォワードに後処理を付け足します。
出力から最小値を引いて最大値と最小値の差分で割れば0~1になり、255をかけると0~255になる、という意味です。

### preprocessを含めて変換

```python
from coremltools.converters.mil.input_types import ColorLayout
ex = torch.randn((1,3,1024,1024)).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=[-0.5,-0.5,-0.5],
                     scale=1/255.0)
    ],
    outputs=[ct.ImageType(name="output",color_layout=ColorLayout.GRAYSCALE)])
coreml_model.save("RMBG.mlpackage")
```

0~255の入力画像を255で割って、0.5を引けば-0.5~0.5に正規化される、という意味です。

# iOSで使う

```swift
guard let model = try? VNCoreMLModel(for: RMBG().model) else {
    fatalError()
}

let request = VNCoreMLRequest(model: model, completionHandler: { [weak self] request, error in
    guard let results = request.results, let firstResult = results.first as? VNPixelBufferObservation else {
        return
    }
    let ciImage = CIImage(cvPixelBuffer: firstResult.pixelBuffer)
    let context = CIContext()
    guard let cgImage = context.createCGImage(ciImage, from: ciImage.extent) else { return }
    DispatchQueue.main.async {
        self?.resultImage = UIImage(cgImage: cgImage)
    }
})

request.imageCropAndScaleOption = .scaleFill

let handler = VNImageRequestHandler(ciImage: ciImage, options: [:])
DispatchQueue.global(qos: .userInteractive).async {
    do {
        try handler.perform([self.request])
    } catch {
        fatalError()
    }
}
```

画像全体を入力にしたいので、
imageCropAndScaleOption は .scaleFill で。
出力は1024*1024の正方形になっているので、元画像のサイズにリサイズして使います。


![IMG_0895.JPG](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/235259/2e1a9741-a2bc-a113-9524-9bbb4e1193a1.jpeg)

変換デモ。

https://colab.research.google.com/drive/1KUTREfLbDklnKhaaJU2Q444TLxjmIXwN?usp=sharing


オリジナルモデル。

https://huggingface.co/briaai/RMBG-1.4

オリジナルモデルのPythonでの使い方はこちら。

https://qiita.com/john-rocky/items/49df2050ec59dc4a008b


🐣
****
フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。

もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…

いずれも中間コストを省いたリーズナブルな価格でお請けできます。

お仕事のご相談はこちらまで
rockyshikoku@gmail.com

機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。

[X](https://twitter.com/JackdeS11)
[Medium](https://medium.com/@rockyshikoku)
[GitHub](https://github.com/john-rocky)