Quickly compute average color of all the clusters generated by connectedComponents

I split an image into clusters with cv2.connectedComponentWithStats. Now I want to observe color changes in each of the clusters.

Currently my code looks like this:

#...Calculation of the mask

res,labels,stats,centroids = cv2.connectedComponentsWithStats(im_mask)

def compute_average(frame,i):
   data=frame[labels==i].mean(axis=0)
   return (data[2]-data[1]) # difference between red and green channel is meaningful for me

while True:
   frame=capture.read()
   if(not frame[0]):
      break
   start_time=time.time()
   measurements = [ compute_average(frame,i) for i in range(1,len(centroids))]
   print("Computing took",time.time()-start_time

It appears that calculating measurements took almost 1.5 second for each frame (I have approximately 300 clusters of 200-600 pixels each). This is unacceptable.

It seems that by cleverly choosing numpy algorithm for computing average I can get much better performance. In particular it should be possible to compute average for all clusters simultaneously. However, I'm stuck here.

Is there a way to sort pixels of image into groups according to their label?

Seems like a perfect use-case to leverage np.bincount which is a pretty efficient way to compute binned summations and weighted summations with its optional second argument. In our case here, we will use the labels as bins and get the summations as the counts and then frame as the weights as the optional weights arg.

Hence, we will have a vectorized and hopefully more efficient way, like so -

def bincount_method(frame, labels):
    f = frame.reshape(-1,3)
    cn1 = np.bincount(labels.ravel(), f[:,1])
    cn2 = np.bincount(labels.ravel(), f[:,2])
    return (cn2[1:]-cn1[1:])/stats[1:,-1]

Benchmarking

For the timings, we will re-use @Dan Mašek's test-setup image. We are using Python 3.6.8.

import numpy as np
import cv2
import urllib.request as ur

# Setup
url = 'https://i.stack.imgur.com/puxMo.png'
s = ur.urlopen(url)
url_response = ur.urlopen(url)
img_array = np.array(bytearray(url_response.read()), dtype=np.uint8)
img = cv2.imdecode(img_array, -1)
im_mask = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

res,labels,stats,centroids = cv2.connectedComponentsWithStats(im_mask)

np.random.seed(0)
frame = np.random.rand(im_mask.shape[0],im_mask.shape[1],3)

Timings -

# Original soln by OP
In [5]: %timeit [compute_average(frame,i) for i in range(1,len(centroids))]
2.38 s ± 116 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# @Dan Mašek's soln
In [3]: %timeit rg_mean_diff_per_label(frame, labels)
92.5 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# Solution with bincount
In [4]: %timeit bincount_method(frame, labels)
30.1 ms ± 82.7 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Allowing pre-computing on labels

Well the bincount one has a small modification for that :

L = (labels.ravel()[:,None]+np.arange(2)*res).ravel()

def bincount_method_with_precompute(frame, L):
    p = np.bincount(L,frame[...,1:].ravel())
    return (p[res+1:]-p[1:res])/stats[1:,-1]

Comparing against @Dan Mašek's solution with pre-compute on the same setup :

In [4]: %timeit bincount_method_with_precompute(frame, L)
25.1 ms ± 326 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [5]: %timeit rg_mean_diff_per_label(frame, label_indices)
20.3 ms ± 432 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

The following algorithm comes to mind. This runs in about 0.115 seconds compared to 2.65 seconds of the original. (For comparison @Divakar's solution runs in about 0.058 seconds)

• Flatten the labels into a linear array. (Note: Better use ravel() since apparently flatten() always makes a copy)
• Perform an indirect sort of the labels to get a list of indices.
• Use this to sort the labels.
• Find the first occurrence of each label in the sorted list (where two neighbor labels differ)
• Then for each label
• Get the range of indices corresponding to this label
• Select the corresponding pixels, calculate mean of columns 1 and 2
• Store the difference

Code:

import numpy as np

def rg_mean_diff_per_label(frame, labels):
    flat_labels = labels.ravel()
    order = flat_labels.argsort()
    sorted_labels = flat_labels[order]
    # Find the position of last occurence of each label
    boundaries = np.where(sorted_labels[:-1] != sorted_labels[1:])[0]
    # And turn it into position of the first occurence each label
    # (except 0, which we want to ignore anyway, as it represents the background)
    boundaries += 1
    # Add position of end of array, so we can simply make ranges using zip(...)
    boundaries = np.append(boundaries, len(flat_labels))

    flat_pixels = frame.reshape(-1, 3) # One pixel per row, 3 columns, 1 per colour channel
   
    measurements = []
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        indices = order[start:end]
        # NB: We don't care about blue, skip it
        data = flat_pixels[indices,1:]
        means = data.mean(axis=0)
        measurements.append(means[1] - means[0])
    
    return measurements

Test image:

Colour data was randomized, the values there don't matter for performance evaluation.

Since you use the same mask for all the frames, i.e. labels remain the same, we can avoid a lot of repeated calculations by refactoring this a little, and caching the parts that stay the same.

Let's look at the profile of the function:

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    88                                           def rg_mean_diff_per_label(frame, labels):
    89         1        109.0    109.0      0.0      flat_labels = labels.ravel()
    90         1     592417.0 592417.0     35.1      order = flat_labels.argsort()
    91         1     107003.0 107003.0      6.3      sorted_labels = flat_labels[order]
    93         1      38591.0  38591.0      2.3      boundaries = np.where(sorted_labels[:-1] != sorted_labels[1:])[0]
    96         1        364.0    364.0      0.0      boundaries += 1
    98         1        666.0    666.0      0.0      boundaries = np.append(boundaries, len(flat_labels))
   100         1         61.0     61.0      0.0      flat_pixels = frame.reshape(-1, 3) # One pixel per row, 3 columns, 1 per colour channel
   102         1         25.0     25.0      0.0      measurements = []
   103       459      11182.0     24.4      0.7      for start, end in zip(boundaries[:-1], boundaries[1:]):
   104       458      17117.0     37.4      1.0          indices = order[start:end]
   106       458     314348.0    686.3     18.6          data = flat_pixels[indices,1:]
   107       458     579712.0   1265.7     34.4          means = data.mean(axis=0)
   108       458      24001.0     52.4      1.4          measurements.append(means[1] - means[0])
   110         1         21.0     21.0      0.0      return measurements

The only thing that changes per iteration is the pixel data, so if we pre-calculate an array of indices of all pixels belonging to each label, we should be able to cut the per-iteration time down by another 40% or so. (This one runs at around 0.057 seconds per iteration but competition has since improved )

Code:

def precalc_label_indices(labels):
    flat_labels = labels.ravel()
    order = flat_labels.argsort()
    sorted_labels = flat_labels[order]
    # Find the position of last occurence of each label
    boundaries = np.where(sorted_labels[:-1] != sorted_labels[1:])[0]
    # And turn it into position of the first occurence each label
    # (except 0, which we want to ignore anyway, as it represents the background)
    boundaries += 1
    # Add position of end of array, so we can simply make ranges using zip(...)
    boundaries = np.append(boundaries, len(flat_labels))
    
    label_indices = []
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        indices = order[start:end]
        indices = np.sort(indices)) # Access in order can't hurt
        label_indices.append(indices)
    return label_indices
    
label_indices = precalc_label_indices(labels)
    
def rg_mean_diff_per_label(frame, label_indices):
    flat_pixels = frame.reshape(-1, 3) # One pixel per row, 3 columns, 1 per colour channel
   
    measurements = []
    for indices in label_indices:
        # NB: We don't care about blue, skip it
        data = flat_pixels[indices,1:]
        means = data.mean(axis=0)
        measurements.append(means[1] - means[0])
    
    return measurements

Now, since you already use OpenCV, why not take advantage of it here? Calculation of the mean seems to be the biggest bottleneck now. As it turns out, the OpenCV mean() is much faster here.

def rg_mean_diff_per_label(frame, label_indices):
    flat_pixels = frame.reshape(-1, 3) # One pixel per row, 3 columns, 1 per colour channel
   
    measurements = []
    for indices in label_indices:
        # NB: We don't care about blue, skip it
        data = flat_pixels[indices,1:]
        means = cv2.mean(data.reshape(-1,1,2)) # This one works per-channel
        measurements.append(means[1] - means[0])
    
    return measurements

This algorithm results in slightly different output (differences on the order of 1e-14 max). However, it now runs in 0.021 seconds per iteration (excluding the pre-calculation which is insignificant in long-run).

Let's look at the profile of the new function.

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    32                                           def rg_mean_diff_per_label(frame, label_indices):
    33        16       3314.0    207.1      0.1      flat_pixels = frame[...,1:].reshape(-1, 2)
    35        16        464.0     29.0      0.0      measurements = []
    36      7344     151147.0     20.6      2.7      for indices in label_indices:
    38      7328    4574161.0    624.2     81.6          data = flat_pixels[indices]
    39      7328     640719.0     87.4     11.4          means = cv2.mean(data.reshape(-1,1,2))
    40      7328     237797.0     32.5      4.2          measurements.append(means[1] - means[0])
    42        16        413.0     25.8      0.0      return measurements

Now the biggest bottleneck is the extraction of pixels that belong to the given label.

The fact that cv2.mean supports a mask parameter gave me another idea. Let's take advantage of the stats, since they contain the bounding box and pixel count for each label. We can pre-generate ROI-sized mask images for each label.
Then for each label, we extraact the ROI (based on the bounding box), and calculate a mean of it with appropriate mask. This time we calculate mean for all three channels, in order to avoid copies of the pixel data.

label_data = []
for label in range(res):
    mask = cv2.inRange(labels, label, label)
    x,y,w,h,n = stats[label]
    roi = mask[y:y+h,x:x+w]
    label_data.append((x,y,x+w,y+h,roi))
    
def test_v4(frame, label_data):
    measurements = []
    for x1,y1,x2,y2,mask in label_data[1:]:
        roi = frame[y1:y2,x1:x2,:]
        means = cv2.mean(roi, mask)
        measurements.append(means[2] - means[1])
    return measurements

This runs in 0.007 seconds per iteration, again excluding the pre-calculation time.

This is using the single test image I created.

Note: I have an idea for a bad-case (3-pixel wide diagonal stripes) scenario where the masks would be very large for the number of pixels contained. I'll provide an update soon.