ハイクオリティな背景除去をiOSで
オリジナルモデルとCoreMLToolsのインストール
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は以下。
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を含んだラップモデルを作る
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を含めて変換
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で使う
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の正方形になっているので、元画像のサイズにリサイズして使います。
変換デモ。
オリジナルモデル。
オリジナルモデルのPythonでの使い方はこちら。
🐣
フリーランスエンジニアです。
AIについて色々記事を書いていますのでよかったらプロフィールを見てみてください。
もし以下のようなご要望をお持ちでしたらお気軽にご相談ください。
AIサービスを開発したい、ビジネスにAIを組み込んで効率化したい、AIを使ったスマホアプリを開発したい、
ARを使ったアプリケーションを作りたい、スマホアプリを作りたいけどどこに相談したらいいかわからない…
いずれも中間コストを省いたリーズナブルな価格でお請けできます。
お仕事のご相談はこちらまで
rockyshikoku@gmail.com
機械学習やAR技術を使ったアプリケーションを作っています。
機械学習／AR関連の情報を発信しています。