Ruslan
Ruslan

Reputation: 122

How can I extract image segment with specific color in OpenCV?

I work with logos and other simple graphics, in which there are no gradients or complex patterns. My task is to extract from the logo segments with letters and other elements.

enter image description here

To do this, I define the background color, and then I go through the picture in order to segment the images. Here is my code for more understanding:

MAXIMUM_COLOR_TRANSITION_DELTA = 100  # 0 - 765


def expand_segment_recursive(image, unexplored_foreground, segment, point, color):
    height, width, _ = image.shape
    # Unpack coordinates from point
    py, px = point

    # Create list of pixels to check
    neighbourhood_pixels = [(py, px + 1), (py, px - 1), (py + 1, px), (py - 1, px)]

    allowed_zone = unexplored_foreground & np.invert(segment)

    for y, x in neighbourhood_pixels:
        # Add pixel to segment if its coordinates within the image shape and its color differs from segment color no
        # more than MAXIMUM_COLOR_TRANSITION_DELTA
        if y in range(height) and x in range(width) and allowed_zone[y, x]:
            color_delta = np.sum(np.abs(image[y, x].astype(np.int) - color.astype(np.int)))
            print(color_delta)
            if color_delta <= MAXIMUM_COLOR_TRANSITION_DELTA:
                segment[y, x] = True
                segment = expand_segment_recursive(image, unexplored_foreground, segment, (y, x), color)
                allowed_zone = unexplored_foreground & np.invert(segment)

    return segment


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Pass image as the argument to use the tool")
        exit(-1)

    IMAGE_FILENAME = sys.argv[1]
    print(IMAGE_FILENAME)

    image = cv.imread(IMAGE_FILENAME)
    height, width, _ = image.shape

    # To filter the background I use median value of the image, as background in most cases takes > 50% of image area.
    background_color = np.median(image, axis=(0, 1))
    print("Background color: ", background_color)

    # Create foreground mask to find segments in it (TODO: Optimize this part)
    foreground = np.zeros(shape=(height, width, 1), dtype=np.bool)
    for y in range(height):
        for x in range(width):
            if not np.array_equal(image[y, x], background_color):
                foreground[y, x] = True

    unexplored_foreground = foreground

    for y in range(height):
        for x in range(width):
            if unexplored_foreground[y, x]:
                segment = np.zeros(foreground.shape, foreground.dtype)
                segment[y, x] = True
                segment = expand_segment_recursive(image, unexplored_foreground, segment, (y, x), image[y, x])

                cv.imshow("segment", segment.astype(np.uint8) * 255)

                while cv.waitKey(0) != 27:
                    continue

Here is the desired result: enter image description here

In the end of run-time I expect 13 extracted separated segments (for this particular image). But instead I got RecursionError: maximum recursion depth exceeded, which is not surprising as expand_segment_recursive() can be called for every pixel of the image. And since even with small image resolution of 600x500 i got at maximum 300K calls.

My question is how can I get rid of recursion in this case and possibly optimize the algorithm with Numpy or OpenCV algorithms?

Upvotes: 1

Views: 6877

Answers (1)

api55
api55

Reputation: 11420

You can actually use a thresholded image (binary) and connectedComponents to do this job in a couple of steps. Also, you may use findContours or other methods.

Here is the code:

import numpy as np
import cv2

# load image as greyscale
img = cv2.imread("hp.png", 0)

# puts 0 to the white (background) and 255 in other places (greyscale value < 250)
_, thresholded = cv2.threshold(img, 250, 255, cv2.THRESH_BINARY_INV)

# gets the labels and the amount of labels, label 0 is the background
amount, labels = cv2.connectedComponents(thresholded)

# lets draw it for visualization purposes
preview = np.zeros((img.shape[0], img.shape[2], 3), dtype=np.uint8)

print (amount) #should be 3 -> two components + background

# draw label 1 blue and label 2 green
preview[labels == 1] = (255, 0, 0)
preview[labels == 2] = (0, 255, 0)

cv2.imshow("frame", preview)
cv2.waitKey(0)

At the end, the thresholded image will look like this:

enter image description here

and the preview image (the one with the colored segments) will look like this:

enter image description here

With the mask you can always use numpy functions to get things like, coordinates of the segments you want or to color them (like I did with preview)

UPDATE

To get different colored segments, you may try to create a "border" between the segments. Since they are plain colors and not gradients, you can try to do an edge detector like canny and then put it black in the image....

import numpy as np
import cv2

img = cv2.imread("total.png", 0)

# background to black
img[img>=200] = 0
# get edges
canny = cv2.Canny(img, 60, 180)
# make them thicker
kernel = np.ones((3,3),np.uint8)
canny = cv2.morphologyEx(canny, cv2.MORPH_DILATE, kernel)
# apply edges as border in the image
img[canny==255] = 0

# same as before
amount, labels = cv2.connectedComponents(img)
preview = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
print (amount) #should be 14 -> 13 components + background

# color them randomly
for i in range(1, amount):
  preview[labels == i] = np.random.randint(0,255, size=3, dtype=np.uint8)

cv2.imshow("frame", preview )
cv2.waitKey(0)

The result is:

enter image description here

Upvotes: 3

Related Questions