import numpy as np
from cv2_41 import cv2
import time

im_mask = cv2.imread('mask.png', cv2.IMREAD_GRAYSCALE)
height, width = im_mask.shape

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

def compute_averages(frame):
    measurements = [ compute_average(frame,i) for i in range(1,len(centroids))]
    return measurements

def test(frame, labels):
    stacked = np.dstack([frame, labels])
    stacked = stacked.reshape(-1, 4)
    sorted_stacked = stacked[stacked[:,3].argsort()]

    indices = sorted_stacked[:,3].flatten()
    boundaries = np.where(indices[:-1] != indices[1:])[0] + 1
    boundaries = np.append(boundaries, len(stacked))

    measurements = []
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        data = sorted_stacked[start:end,:3].mean(axis=0)
        measurements.append(data[2]-data[1])

    return measurements

while True:
    frame = np.zeros((height, width, 3), np.uint8)
    cv2.randu(frame, (0,0,0), (256,256,256))
    start_time=time.time()
    measurements_a = compute_averages(frame)
    print("Computing A took",time.time()-start_time)

    start_time=time.time()
    measurements_b = test(frame, labels)
    print("Computing B took",time.time()-start_time)

    print(np.all(np.array(measurements_a) == np.array(measurements_b)))