Reusing models from grabcut in OpenCV
I used the interactive grabcut.py from the OpenCV samples to segment the image and kept the foreground and background models. Then I used these models to segment more images of the same type, since I don't want to retrain the model every time.
After running the grabcut algorithm, the mask is all zeros (all background) and therefore it does not segment anything.
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()
I tried to initialize the algorithm with a mask or rectangle, but that throws an error because the models are not empty (which is what I really want).
How do I feed my trained models to the algorithm so they don't retrain from scratch every time I segment an image?
EDIT Following Rayryeng's comment, I ran the following code:
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 0, cv2.GC_INIT_WITH_RECT)
cv2.grabCut(img, mask, rect, bgdModel, fgdModel, 2, cv2.GC_EVAL)
Seems to work, but the first call now changes my model. In the source code, it calls learnGMMs
without checking for a preset model.
source to share
You have the right line of thinking that you are using cv2.GC_EVAL
so that you need to perform segmentation without having to compute the models again. Unfortunately, even if you use this flag, it is a limitation with the OpenCV source itself. If you look at the actual C ++ implementation, when you encounter a conditionGC_EVAL
, it does so at the end of the methodcv::grabcut
. Note that the Python method cv2.grabCut
is a wrapper for 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 );
}
You will see that it GC_EVAL
appears only once in the code and validates the input. The culprit is function learnGMMs
. Even if you specified trained models, they get reset because the call learnGMMs
ignores the flag GC_EVAL
, so it gets called regardless of which flag you specify as input.
Inspired by this post : OpenCV - GrabCut with custom foreground / background models , what you can do you have to change the OpenCV source yourself and inside the loop you can put an instruction if
to check the flag GC_EVAL
before calling learnGMMs
:
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 );
}
It should be able to use pre-trained models without having to learn them over and over at each iteration. After making the changes, you will have to recompile the source again and hopefully we can use your pretrained models without clearing them when you use the flag cv2.GC_EVAL
.
In the future, I opened a question about the official repo for OpenCV. Hopefully they fix it when they have time: https://github.com/opencv/opencv/issues/9191
source to share