在 OpenCV 中重用 grabcut 模型

Reusing models from grabcut in OpenCV

我使用 OpenCV 示例中的交互式 grabcut.py 来分割图像并保存前景和背景模型。然后我使用这些模型来分割更多同类图像,因为我不想每次都重新训练模型。

在 运行 抓取算法之后,掩码全为零(所有背景),因此它不会分割任何东西。

from matplotlib import pyplot as plt
import numpy as np
import cv2

img = cv2.imread('usimg1.jpg')
mask = np.zeros(img.shape[:2], np.uint8)
bgdModel = np.load('bgdmodel.npy')
fgdModel = np.load('fgdmodel.npy')

cv2.grabCut(img, mask, None, bgdModel, fgdModel, 5, cv2.GC_EVAL)

mask = np.where((mask==2) | (mask==0), 0, 1).astype('uint8') 
img = img * mask[:, :, np.newaxis]

plt.imshow(img)
plt.show()

我尝试用掩码或矩形初始化算法,但这会产生错误,因为模型不为空(这正是我真正想要的)。

我如何才能将预训练模型传递给算法,这样它们就不会在我每次分割图像时都从头开始重新训练?

编辑 在 rayryeng 的评论之后,我实现了以下代码:

cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 0, cv2.GC_INIT_WITH_RECT)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 2, cv2.GC_EVAL)

它似乎有效,但现在第一个调用改变了我的模型。在源代码中它调用 learnGMMs 而不检查是否提供了预训练模型。

您在使用 cv2.GC_EVAL 时的思路是正确的,因此您只需执行分割而无需再次计算模型。 不幸的是,即使您使用此标志,这也是 OpenCV 源代码本身的限制。如果你看actual C++ implementation when you encounter the GC_EVAL condition, it does this towards the end of the cv::grabcut method。请注意,Python cv2.grabCut 方法是 cv::grabcut:

的包装器
if( mode == GC_EVAL )
    checkMask( img, mask );

const double gamma = 50;
const double lambda = 9*gamma;
const double beta = calcBeta( img );

Mat leftW, upleftW, upW, uprightW;
calcNWeights( img, leftW, upleftW, upW, uprightW, beta, gamma );

for( int i = 0; i < iterCount; i++ )
{
    GCGraph<double> graph;
    assignGMMsComponents( img, mask, bgdGMM, fgdGMM, compIdxs );
    learnGMMs( img, mask, compIdxs, bgdGMM, fgdGMM );
    constructGCGraph(img, mask, bgdGMM, fgdGMM, lambda, leftW, upleftW, upW, uprightW, graph );
    estimateSegmentation( graph, mask );
}

您会看到 GC_EVAL 在代码中只出现一次,那是为了检查输入的有效性。罪魁祸首是 learnGMMs 函数。即使您指定了经过训练的模型,这些模型也会被重置,因为对 learnGMMs 的调用会忽略 GC_EVAL 标志,因此无论您指定什么标志作为输入,都会调用它。

受此 post 的启发:OpenCV - GrabCut with custom foreground/background models,您可以做的是您必须自己修改 OpenCV 源代码,并在循环内放置一个 if 语句来在调用 learnGMMs:

之前检查 GC_EVAL 标志
if( mode == GC_EVAL )
    checkMask( img, mask );

const double gamma = 50;
const double lambda = 9*gamma;
const double beta = calcBeta( img );

Mat leftW, upleftW, upW, uprightW;
calcNWeights( img, leftW, upleftW, upW, uprightW, beta, gamma );

for( int i = 0; i < iterCount; i++ )
{
    GCGraph<double> graph;
    assignGMMsComponents( img, mask, bgdGMM, fgdGMM, compIdxs );
    if (mode != GC_EVAL) // New
        learnGMMs( img, mask, compIdxs, bgdGMM, fgdGMM );
    constructGCGraph(img, mask, bgdGMM, fgdGMM, lambda, leftW, upleftW, upW, uprightW, graph );
    estimateSegmentation( graph, mask );
}

这应该能够使用预训练模型,而无需在每次迭代时重新学习它们。进行更改后,您将不得不再次重新编译源代码,并且当您使用 cv2.GC_EVAL 标志时,应该能够使用您预训练的模型而不清除它们。

为了将来,我已经在 OpenCV 的官方回购上提出了一个问题。希望他们有空的时候会解决这个问题:https://github.com/opencv/opencv/issues/9191