Reusing models from grabcut in OpenCV

Viewed 1668

I used the interactive grabcut.py from the OpenCV samples to segment an image and saved the foreground and background models. Then I used these models to segment more images of the same kind, as I don't want to retrain the model each time.

After running the grabcut algorithm, the mask is all zeros (all background) and therefore it doesn't 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 a rectangle but this produces an error because the models are not empty (which is what I actually want).

How do I have to pass the pre-trained models to the algorithm, such that they are not retrained from scratch each time I'm segmenting an image?

EDIT After rayryeng's comment I implemented 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)

It seems to work but the first call now changes my model. In the source code it calls learnGMMs without checking whether a pretrained model is provided.

1 Answers
Related