from cv2_41 import cv2
import os
import numpy as np
import time

t = time.time()

img = cv2.imread('test.jpg')
print "Initial img shape:", img.shape

for _ in range(0, 3):
    img = cv2.pyrDown(img)

print "Reduced img shape:", img.shape

img_super = cv2.ximgproc.createSuperpixelSLIC(img, cv2.ximgproc.MSLIC, 100, 10.0)
img_super.iterate(3)

labels = img_super.getLabels() # No need to pre-allocate here

superpixel_count = img_super.getNumberOfSuperpixels();
print "Superpixel count:", superpixel_count

super_pixelized = np.zeros_like(img)

mask = np.zeros(img.shape[:2], np.uint8) # Here it seems to make more sense to pre-alloc and reuse
for label in range(0, superpixel_count):
    cv2.inRange(labels, label, label, dst=mask)

    # Find the bounding box of this label
    x,y,w,h = cv2.boundingRect(mask)

    # Work only on the rectangular region containing the label
    mask_roi = mask[y:y+h,x:x+w]
    img_roi = img[y:y+h,x:x+w]

    # Per-channel mean of the masked pixels (we're usingo BGR, so we drop the useless 4th channel it gives us)
    roi_mean = cv2.mean(img_roi, mask_roi)[:3]

    # Set all masked pixels in the ROI of the target image the the mean colour
    super_pixelized[y:y+h,x:x+w][mask_roi != 0] = roi_mean

cv2.imwrite("superpix.png", super_pixelized)

print "TOTAL TIME:", (time.time() - t)