Extracting prominent color gradients from an image

colors.py

import imageio
import numpy as np
import cv2

q = "20191230_221836_HDR.jpg"
r = "legs.png"
im = np.array(imageio.imread(r))
im_orig = im
dims = im.shape

blurfactor = 5
#scale = 0.125
scale = 120/dims[0]
kernel = np.ones((blurfactor,blurfactor),np.float32)/(blurfactor**2)
im_ = im
im_ = cv2.filter2D(im_,-1,kernel)
im_ = cv2.resize(im_, dsize=(int(dims[1]*scale), int(dims[0]*scale)), interpolation=cv2.INTER_CUBIC)

imageio.imwrite("pre.png", im_)

im_rgb = im_
im_hsv = cv2.cvtColor(im_, cv2.COLOR_RGB2HSV)
im = im_rgb
dims = im.shape

im_vecs = np.zeros((dims[0], dims[1], 1+2), dtype=np.uint8)
print (dims)

# Calculate direction of least change for each pixel
for x in range(0, dims[0]-1):
    print (x)
    for y in range(0, dims[1]-1):
        val = im[x,y]

        if (x == 0 and y == 0):
            val_upleft = 2 * val
        else:
            val_upleft = im[x-1,y-1]

        if (x == 0):
            val_left = 2 * val
            val_downleft = 2 * val
        else:
            val_left = im[x-1,y]
            val_downleft = im[x-1,y+1]

        if (y == 0):
            val_up = 2 * val
            val_upright = 2 * val
        else:
            val_up = im[x,y-1]
            val_upright = im[x+1,y-1]

        
        val_down = im[x,y+1]
        val_downright = im[x+1,y+1]
        val_right = im[x+1,y]
        
        alldirs_direction = [(-1,0), (-1,-1), (0,-1), (1,-1), (1,0), (1,1), (0,1), (-1,1)]
        alldirs = [val_left, val_upleft, val_up, val_upright, val_right, val_downright, val_down, val_downleft]
        mindir = -1
        minval = 256
        for i in range(0, 8):
            diff = abs(alldirs[i] - val)
            #if ( <= minval):
            #    mindir = i
            #    minval = np.linalg.norm(diff)
            diff = 0.2126*diff[0] + 0.7152*diff[1] + 0.0722*diff[2]
            #diff = (diff[0] + diff[1] + diff[2]) / 3
            #diff = np.linalg.norm(diff)
            if (diff <= minval):
                mindir = i
                minval = diff

        diff = np.array([minval])#np.array((alldirs[mindir] - val))
        direction = np.array(alldirs_direction[mindir])
        entry = np.concatenate((diff, direction))
        im_vecs[x,y] = entry

# Use direction of least change to find color gradients
visited = np.zeros((dims[0], dims[1]), dtype=np.uint8)
color_visited = np.zeros((256,256,256), dtype=np.uint8)
gradients = np.zeros((dims[0], dims[1], dims[2]), dtype=np.uint8)

gradients_list = []
longest_chain = 0
im = im_rgb
for x in range(0, dims[0]-1):
    for y in range(0, dims[1]-1):
        xpos = x
        ypos = y
        gradient = []
        lastval = np.array([0,0,0])
        lastval_hsv = np.array([0,0,0])
        has_lastval = False
        for i in range(0, 4000):
            if (ypos >= dims[1] or xpos >= dims[0] or xpos < 0 or ypos < 0):
                break
            if (visited[xpos,ypos] == 255):
                break

            #visited[xpos,ypos] = 255
            val = im[xpos,ypos,:]
            val_hsv = im[xpos,ypos,:]
            vec = im_vecs[xpos,ypos,:]
            
            if (len(gradient) > 4):
                prev_var = np.linalg.norm(np.var(gradient, axis=0))
                next_var = np.linalg.norm(np.var(gradient + [val], axis=0))
                if (abs(prev_var - next_var) > 300):
                    break
            if (has_lastval and abs(np.linalg.norm(val_hsv - lastval_hsv)) > 70):
                break
            
            if (color_visited[val[0], val[1], val[2]] != 1):
            #if (True):
                gradient += [val]
                color_visited[val[0], val[1], val[2]] = 1
                gradients[xpos,ypos] = val
                visited[xpos,ypos] = 255
            has_lastval = True
            lastval = val
            lastval_hsv = val_hsv
            #else:
                #gradients[xpos,ypos] = lastval
            
            xpos += vec[1]
            ypos += vec[2]
            
            #print (val, xpos, ypos, vec)

        accepted = True
        if (len(gradient) <= 4):
            accepted = False
        elif (np.linalg.norm(np.var(gradient, axis=0)) < 400):
            accepted = False

        if (accepted):
            gradients_list += [gradient]
            if (len(gradient) > longest_chain):
                longest_chain = len(gradient)
        else:
            for val in gradient:
                color_visited[val[0], val[1], val[2]] = 0

# Generate gradient image
g_x = 0
g_y = 0
gradients_list.sort(key=len, reverse=True)
gradients_list_img = np.zeros((len(gradients_list)+1, longest_chain+1, dims[2]), dtype=np.uint8)
tossed = np.zeros((len(gradients_list),))
palette_set = []
for i in range(0, len(gradients_list)):
    gradient = gradients_list[i]
    
    discard = False
    for j in range(0, i):
        if (tossed[i] == 1):
            continue

        orig_var = np.linalg.norm(np.var(gradient, axis=0))
        comb_var = np.linalg.norm(np.var(gradient+gradients_list[j], axis=0))
        #print (i,j,orig_var,comb_var,abs(comb_var-orig_var))
        if (abs(comb_var-orig_var) < 100):
            tossed[i] = 1
            discard = True
            break
    if (discard):
        continue

    print (np.var(gradient, axis=0), np.linalg.norm(np.var(gradient, axis=0)))
    g_x = 0
    for val in gradient:
        palette_set += [val]
        gradients_list_img[g_y, g_x] = val
        g_x += 1
    g_y += 1

gradients_list_img = cv2.resize(gradients_list_img, dsize=(gradients_list_img.shape[1] * 32, gradients_list_img.shape[0] * 32), interpolation=cv2.INTER_NEAREST)

imageio.imwrite("testi.png", im_vecs[:,:,:])
imageio.imwrite("testj.png", visited)
imageio.imwrite("testk.png", gradients)
imageio.imwrite("testl.png", gradients_list_img)