Mask R CNN, image segmentation etc are all algorithms that have become extremely useful in today’s world. These algorithms perform well because of the concept involving separation between the foreground and the background. Doing this is quite simple when it comes to OpenCV and this process can be done interactively by drawing the outline ourselves.
In this article, we will use an algorithm called GrabCut to interactively segregate the foreground from the background in an image.
Working of the GrabCut algorithm
- The user either chooses to manually select the region of interest or adjusts the mask values in the algorithm to automatically do so. Since we are doing it interactively, we will manually draw the regions. Once the region is selected the area outside the region of interest will be turned black.
- A gaussian mixture model (GMM) is applied over the image and this model understands the user input and starts creating labels for unknown pixel values. These pixels are then grouped into one based on the colour statistics.
- Based on the above pixel distribution, a graph is generated where these pixels are considered as nodes. Apart from this, there are two other nodes called Source and Sink node that is generated.
- All the foreground pixels will be connected to the source node and all the background pixels with the sink node. The edge connecting the source node with the sink node contains the weights which is basically the probability of whether the pixel is a foreground one or background one.
- Now the algorithm segments the image into two with the help of a cost function and separates the source and sink node into two. After multiple iterations of this, the algorithm can finally extract the foreground part of the image and make the background black.
All of these are combined into cv2.grabcut() function in OpenCV.
Implementation of the Algorithm
Let us start by selecting the input image of our choice. I have selected the following image. You can download the image here.
Loading the Image and Libraries
Now, I will load this image and import the necessary libraries. I will also specify the foreground and background of the image using the NumPy array. The array consists of only zeros and is constructed with a dimension of 1 row and 65 columns.
import numpy as np import cv2 input_image = cv2.imread(dino.jpg') background_array = np.zeros((1,65),np.float64) foreground_array = np.zeros((1,65),np.float64)
The first step is to create a mask to separate the foreground and background from the image. But, while displaying the image we still need to display the whole image as one. Hence we will write a function for this and return the individual masks.
def get_merge_mask(foreground_mask, background_mask, input_image): get_mask = foreground_mask[:,:,1:2]/255 + background_mask[:,:,2:3]/255 if np.max(np.max(np.max(get_mask)))>1: return False, input_image get_mask = get_mask.astype('uint8') return True, foreground_mask + background_mask + input_image*(1-get_mask)
Creating a Mouse Event for Drawing
Since we have to manually draw the outline of the part we want to extract, we need to create an interactive mouse event to understand when the mouse has clicked and the location of where it is moving.
First, let us define the parameters that are needed for the event.
draw_lines = False method = True x_value,y_value = -1,-1 r = 8
Now, we will use the OpenCV event listeners for mouse button down for both left and right clicks and detect the mouse movement.
def interactive_draw(event,x_axis,y_axis,flags,param): if event == cv2.EVENT_LBUTTONDOWN: draw_lines = True x_value,y_value = x_axis,y_axis if event == cv2.EVENT_RBUTTONUP: method = not method elif event == cv2.EVENT_MOUSEMOVE: if draw_lines is True: if method is True: cv2.circle(foreground_mask,(x_axis,y_axis),r,(0,255,0),-1) else: cv2.circle(background_mask,(x_axis,y_axis),r,(0,0,255),-1) elif event == cv2.EVENT_LBUTTONUP: draw_lines = False if method is True: cv2.circle(foreground_mask,(x_axis,y_axis),r,(0,255,0),-1) else: cv2.circle(background_mask,(x_axis,y_axis),r,(0,0,255),-1)
Initialize and Define the Method
Now that we have written the functions we will now initialize the required values and define the method as needed.
foreground_mask = np.zeros_like(input_image) background_mask = np.zeros_like(input_image) cv2.namedWindow('grabcut algorithm') cv2.setMouseCallback('grabcut',interactive_draw)
We have initialized the masks and called the mouse callback function. Now we can call the functions and our GrabCut algorithm on the image as follows.
while True: source, apply_mask = get_merge_mask(foreground_mask, background_mask, input_image) if not source: foreground_mask = np.zeros_like(input_image) background_mask = np.zeros_like(input_image) cv2.imshow('image',apply_mask) key = cv2.waitKey(1) & 0xFF if key == ord('m'): method = not method if method: print('front') else: print('back') elif key == 13: print('cropping') output_mask = np.zeros(input_image.shape[:2],np.uint8) + 2 output_mask[foreground_mask[:,:,1] == 255] = 1 output_mask[background_mask[:,:,2] == 255] = 0
Now we will call the GrabCut algorithm to make the cut for us
output_mask, background_array,foreground_array = cv2.grabCut(input_image,output_mask,None,background_array,foreground_array,iteration,cv2.GC_INIT_WITH_MASK) output_mask = np.where((output_mask==2)|(output_mask==0),0,1).astype('uint8') final_out = input_image*output_mask[:,:,np.newaxis] cv2.imshow('final_out',final_out) cv2.imwrite('output.jpg',final_out) elif key == 27: break cv2.destroyAllWindows()
When you run the program the image will pop up and you will have to draw the foreground on the screen.
Then, press ‘enter’ key and you will see that clearly the background is removed and the image is neatly segmented.
In the above article, we saw how to use the GrabCut algorithm to extract foreground from images interactively. These can be very helpful for image segmentation and to remove unwanted noise from images while using them for algorithms.
The complete code of this implementation is available on the AIM’s GitHub repository. Please visit this link for the notebook with complete code.
If you loved this story, do join our Telegram Community.
Also, you can write for us and be one of the 500+ experts who have contributed stories at AIM. Share your nominations here.
What's Your Reaction?
I am an aspiring data scientist with a passion for teaching. I am a computer science graduate from Dayananda Sagar Institute. I have experience in building models in deep learning and reinforcement learning. My goal is to use AI in the field of education to make learning meaningful for everyone.