superpixels using SLIC

SLICcv.py

import cv2
import sys
import numpy as np
import time

class SLIC:
    def __init__(self, img, step, nc):
        self.img = img
        self.height, self.width = img.shape[:2]
        self.labimg = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
        self.step = step
        self.nc = nc
        self.ns = step
        self.FLT_MAX = 1000000
        self.ITERATIONS = 10
        #self._initData()

    def generateSuperPixels(self):
        self._initData()
        print self.centers.shape, "number of centers"
        print self.clusters.shape, "number of clusters"
        print self.center_counts.shape, "number of center_counts"
        print self.centers
        #raw_input("centers:")

        indnp = np.mgrid[0:self.height,0:self.width].swapaxes(0,2).swapaxes(0,1)
        for i in range(self.ITERATIONS):
            self.distances = self.FLT_MAX * np.ones(self.img.shape[:2])
            for j in xrange(self.centers.shape[0]):
                """
                for k in xrange(self.centers[j][3]-self.step, self.centers[j][3]+self.step):
                    for l in xrange(self.centers[j][4]-self.step, self.centers[j][4]+self.step):
                        print j, k, l
                        if k >=0 and k < self.width and l >=0 and l < self.height:
                            color = self.labimg[l, k]
                            d = self._computeDist(j, (k, l), color)
                            print d
                            if d < self.distances[k][l]:
                                self.distances[k][l] = d
                                self.clusters[k][l] = j
                """
                xlow, xhigh = int(self.centers[j][3] - self.step), int(self.centers[j][3] + self.step)
                ylow, yhigh = int(self.centers[j][4] - self.step), int(self.centers[j][4] + self.step)

                if xlow <= 0:
                    xlow = 0
                if xhigh > self.width:
                    xhigh = self.width
                if ylow <=0:
                    ylow = 0
                if yhigh > self.height:
                    yhigh = self.height

                cropimg = self.labimg[ylow : yhigh , xlow : xhigh].astype(np.int64)
                colordiff = cropimg - self.labimg[self.centers[j][4], self.centers[j][3]]
                colorDist = np.sqrt(np.sum(np.square(colordiff.astype(np.int64)), axis=2))

                yy, xx = np.ogrid[ylow : yhigh, xlow : xhigh]
                pixdist = ((yy-self.centers[j][4])**2 + (xx-self.centers[j][3])**2)**0.5
                dist = ((colorDist/self.nc)**2 + (pixdist/self.ns)**2)**0.5

                distanceCrop = self.distances[ylow : yhigh, xlow : xhigh]
                idx = dist < distanceCrop
                distanceCrop[idx] = dist[idx]
                self.distances[ylow : yhigh, xlow : xhigh] = distanceCrop
                self.clusters[ylow : yhigh, xlow : xhigh][idx] = j

            # clear the center values
            # save some more time by not clearing it
            #print "clear the center values"
            #self.centers = np.zeros((self.centers.shape[0], 5))
            #self.center_counts = np.zeros(self.centers.shape[0])

            # compute new cluster values
            print "compute new cluster values"
            t = time.time()
            for k in xrange(len(self.centers)):
                idx = (self.clusters == k)
                colornp = self.labimg[idx]
                distnp = indnp[idx]
                self.centers[k][0:3] = np.sum(colornp, axis=0)
                sumy, sumx = np.sum(distnp, axis=0)
                self.centers[k][3:] = sumx, sumy
                #self.centers[k][0:3] = np.average(colornp, axis=0)
                #self.centers[k][3:] = np.average(distnp, axis=0)
                # Normalize here
                self.centers[k] /= np.sum(idx)
                #self.center_counts[k] = np.sum(idx)
            #print self.centers
            t11 = time.time()
            print t11 - t, "time taken for clusters"
            """
            for j in xrange(self.width):
                for k in xrange(self.height):
                    c_id = self.clusters[k, j]
                    #print c_id, "c_id"
                    if c_id != -1:
                        color = self.labimg[k, j]
                        self.centers[c_id][0:3] += color
                        self.centers[c_id][3:] += (j, k)
                        self.center_counts[c_id] +=1
            """
            # Normalize the cluster
            """
            print "Normalize the cluster"
            for j in xrange(self.centers.shape[0]):
                self.centers[j] /= self.center_counts[j]
            """
        print self.centers

    def _initData(self):
        print "_initData"
        print self.step, "step"
        self.clusters = -1 * np.ones(self.img.shape[:2])
        self.distances = self.FLT_MAX * np.ones(self.img.shape[:2])

        centers = []
        for i in xrange(self.step, self.width - self.step/2, self.step):
            for j in xrange(self.step, self.height - self.step/2, self.step):
                #print i, j, "_findLocalMinimum"
                nc = self._findLocalMinimum(center=(i, j))
                color = self.labimg[nc[1], nc[0]]
                center = [color[0], color[1], color[2], nc[0], nc[1]]
                centers.append(center)
        self.center_counts = np.zeros(len(centers))
        self.centers = np.array(centers)
        print self.centers.shape, "number of centers"
        print "_intiData finished"

    def createConnectivity(self):
        print "createConnectivity"
        label = 0
        adjlabel = 0
        lims = self.width * self.height / self.centers.shape[0]
        dx4 = [-1, 0, 1, 0]
        dy4 = [0, -1, 0, 1]
        new_clusters = -1 * np.ones(self.img.shape[:2]).astype(np.int64)
        print lims, "lims"
        #idx = new_clusters == -1
        elements = []
        for i in xrange(self.width):
            for j in xrange(self.height):
                #print new_clusters[j, i]
                #print j, i
                if new_clusters[j, i] == -1:
                    elements = []
                    elements.append((j, i))
                    for dx, dy in zip(dx4, dy4):
                        x = elements[0][1] + dx
                        y = elements[0][0] + dy
                        if (x>=0 and x < self.width and 
                            y>=0 and y < self.height and 
                            new_clusters[y, x] >=0):
                            adjlabel = new_clusters[y, x]
                #print elements
                count = 1
                """
                for dx, dy in zip(dx4, dy4):
                    x = elements[0][1] + dx
                    y = elements[0][0] + dy

                    if (x>=0 and x<self.width and y>=0 and y<self.height):
                        if new_clusters[y, x] == -1 and self.clusters[j, i] == self.clusters[y, x]:
                            elements.append((y, x))
                            new_clusters[y, x] = label
                            count+=1
                """

                
                #for c in range(count):
                c = 0
                while c < count:
                    for dx, dy in zip(dx4, dy4):
                        x = elements[c][1] + dx
                        y = elements[c][0] + dy

                        if (x>=0 and x<self.width and y>=0 and y<self.height):
                            if new_clusters[y, x] == -1 and self.clusters[j, i] == self.clusters[y, x]:
                                elements.append((y, x))
                                new_clusters[y, x] = label
                                count+=1
                    c+=1
                #print count
                if (count <= lims >> 2):
                    for c in range(count):
                        new_clusters[elements[c]] = adjlabel
                    label-=1
                label+=1
        print label
        self.new_clusters = new_clusters

    def displayContours(self, color):
        print "displayContours"
        dx8 = [-1, -1, 0, 1, 1, 1, 0, -1]
        dy8 = [0, -1, -1, -1, 0, 1, 1, 1]

        isTaken = np.zeros(self.img.shape[:2], np.bool)
        contours = []

        for i in xrange(self.width):
            for j in xrange(self.height):
                nr_p = 0
                for dx, dy in zip(dx8, dy8):
                    x = i + dx
                    y = j + dy
                    if x>=0 and x < self.width and y>=0 and y < self.height:
                        if isTaken[y, x] == False and self.clusters[j, i] != self.clusters[y, x]:
                            #print self.clusters[j, i]
                            #print self.clusters[y, x]
                            nr_p += 1
                #print nr_p, "nr_p", j, i
                if nr_p >= 2:
                    isTaken[j, i] = True
                    contours.append([j, i])
        #print contours
        #print len(contours)
        for i in xrange(len(contours)):
            self.img[contours[i][0], contours[i][1]] = color


    def _findLocalMinimum(self, center):
        min_grad = self.FLT_MAX
        loc_min = center
        #print "_findLocalMinimum"
        #print center
        for i in xrange(center[0] - 1, center[0] + 2):
            for j in xrange(center[1] - 1, center[1] + 2):
                #print i, j
                c1 = self.labimg[j+1, i]
                c2 = self.labimg[j, i+1]
                c3 = self.labimg[j, i]
                #print ((c1[0] - c3[0])**2)**0.5
                #print ((c2[0] - c3[0])**2)**0.5
                if ((c1[0] - c3[0])**2)**0.5 + ((c2[0] - c3[0])**2)**0.5 < min_grad:
                    min_grad = abs(c1[0] - c3[0]) + abs(c2[0] - c3[0])
                    loc_min = [i, j]
        #print loc_min, "loc_min"
        return loc_min

    def _computeDist(self, ci, pixel, color):
        dc = (sum((self.centers[ci][:3] - color)**2))**0.5
        ds = (sum((self.centers[ci][3:] - pixel)**2))**0.5
        return ((dc/self.nc)**2 + (ds/self.ns)**2)**0.5

img = cv2.imread(sys.argv[1])
nr_superpixels = int(sys.argv[2])
nc = int(sys.argv[3])
#img = cv2.resize(img, (300, 200))
step = int((img.shape[0]*img.shape[1]/nr_superpixels)**0.5)
print step, "step"
slic = SLIC(img, step, nc)
slic.generateSuperPixels()
slic.createConnectivity()
cv2.imshow("cluster", slic.clusters.astype(np.uint8))
print slic.clusters
cv2.imshow("img", slic.img)
slic.displayContours((255, 0, 0))
cv2.imshow("fin", slic.img)
cv2.waitKey(0)
cv2.imwrite("SLICimg.jpg", slic.img)