画像内のオブジェクトを綺麗に切り抜く

画像を入力するだけで、画像内の目立つオブジェクトを綺麗に切り抜いて、セグメンテーションしてくれるモデルがあります。
型抜きのように、隅々まで綺麗に分離できれば、画像合成などさまざまな用途に使えます。

最新の高精度モデルIS-Net

過去には、U2-netという、かなり高精度のセマンティックセグメンテーションモデルがあり、さまざまなアプリケーションに使われています。
同じ作者による高精度モデルが２０２２年７月に発表されました。
IS-Netです。

(公式リポジトリより引用)

使い方

Python

推論Colab

以下のColabノートブックを実行すると、結果がresultsフォルダに保存されます。

もしくは、リポジトリをクローンして、事前トレーニング済みモデルをダウンロード後、以下のスクリプトで推論できます。

import os
import numpy as np
from skimage import io

import glob

import torch, gc
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import torch.nn.functional as F
from torchvision.transforms.functional import normalize

from data_loader_cache import get_im_gt_name_dict, create_dataloaders, GOSRandomHFlip, GOSResize, GOSRandomCrop, GOSNormalize #GOSDatasetCache,
from basics import  f1_mae_torch #normPRED, GOSPRF1ScoresCache,f1score_torch,
from models import *

net = ISNetDIS() 
net.load_state_dict(torch.load("isnet.pth",map_location="cpu"))
net.eval()
net.cuda()

im = io.imread(image_file)/255
w,h,_ = im.shape
if len(im.shape) < 3:
    im = im[:, :, np.newaxis]
if im.shape[2] == 1:
    im = np.repeat(im, 3, axis=2)
im_tensor = torch.tensor(im, dtype=torch.float32).cuda()
im_tensor = torch.transpose(torch.transpose(im_tensor,1,2),0,1)
im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1,1,1])
im_tensor = torch.unsqueeze(im_tensor,0)
im_tensor = F.interpolate(im_tensor, size=(1024,1024))

ds_val = net(im_tensor)[0]
im_pred = F.interpolate(ds_val[0], size=(w,h))

im_pred = torch.squeeze(im_pred)
ma = torch.max(im_pred)
mi = torch.min(im_pred)
im_pred = (im_pred-mi)/(ma-mi)
im_result = im_pred.to('cpu').detach().numpy().copy()
io.imsave(os.path.join(result_folder,'output.jpg'), im_result)

Swift

iOSとMacOSではCoreMLに変換したモデルが使えます。

推論リクエストを実行します。

guard let model = try? DIS_1024_2048(configuration: MLModelConfiguration()).model,
      let vnModel = try? VNCoreMLModel(for: model) else {
   fatalError()
}
let request = VNCoreMLRequest(model: vnModel,completionHandler: mlCompletionHandler(request:error:))
request.imageCropAndScaleOption = .scaleFill
 
let handler = VNImageRequestHandler(ciImage: ciImage)
do {
    try handler.perform([request])
    guard let result = request?.results?.first as? VNPixelBufferObservation else {
        fatalError()
    }
    let pixelBuffer = result.pixelBuffer
} catch let error {
    fatalError(error.localizedDescription)
}

🐣

フリーランスエンジニアです。
お仕事のご相談こちらまで
rockyshikoku@gmail.com

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium
GitHub