Search code examples
pythonimage-segmentationmask-rcnn

How to extract a segmented object or change the background of the original image?


I am doing a detection with Mask R-CNN of one model available at Train Mask R-CNN for Image Segmentation.

  • Code I
# Load Image
img = cv2.imread("/content/image.jpg")

test_model, inference_config = load_inference_model(1, "/content/mask_rcnn.h5")
image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# Detect results
r = test_model.detect([image])[0]
colors = random_colors(80)

# Get Coordinates and show it on the image
object_count = len(r["class_ids"])
for i in range(object_count):
    # 1. Mask
    mask = r["masks"][:, :, i]
    contours = get_mask_contours(mask)
    for cnt in contours:
        cv2.polylines(img, [cnt], True, colors[i], 2)
        img = draw_mask(img, [cnt], colors[i])
  • First attempt

After training I tried to change the background color of the image using the code available at Change the Background of Any Image with 5 Lines of Code. Although it is working, it is not worth it for me to use another model to change the background.

  • Second attempt

Using a suggestion that is also available at Can anyone tell me how can I change mask the foreground if I know the color range of background in RGB?. I added the following line of code img2[ ~mask ] = [0,0,0].

This solution works if for the detected object is only one mask, otherwise it will not work so well. Example:

Example

To solve this problem I created an empty list and did the operation inside of the for loop, but it didn't work.

I noticed that in the current way, only the last mask is taken. Is there a way to get all masks? because if there is more than one object in the image this code doesn't work either.


Solution

  • I can't test it but if you have mask as numpy.array with the same size as image then you can use it to replace pixels in image

    image[ ~mask ] = [0,0,0]
    

    For example

    python - Can anyone tell me how can I change mask the foreground if I know the color range of background in RGB? - Stack Overflow

    Finding red color in image using Python & OpenCV - Stack Overflow


    If you have more masks then you can use operation OR on all masks to create one mask

    final_mask = mask1 | mask2 | mask3 
    

    In your code it could be

    final_mask = r["masks"][:, :, 0] | r["masks"][:, :, 1] | r["masks"][:, :, 2]
    

    or using loop (to make it more universal)

    final_mask = r["masks"][:, :, 0] 
    
    for i in range(1, object_count): 
        final_mask |= r["masks"][:, :, i]
    

    If you would have list of masks then you could use

    final_mask = masks[0] | masks[1] | masks[2]
    

    or you could use function functools.reduce() (to make it more universal)

    from functools import reduce
    
    def or_masks(a, b):
        return a | b
    
    final_mask = reduce(or_masks, masks)
    

    or shorter

    from functools import reduce
    
    final_mask = reduce(lambda a,b:a|b, masks)