Image theme color extraction algorithm

Keywords: Python Algorithm AI

Median segmentation

In RGB color space, the three primary colors R, G and B correspond to three coordinate axes of the space, and each coordinate axis is quantized to 0-255. 0 corresponds to all black and 255 corresponds to all white. In this way, a color cube with 256 sides is formed, and all possible colors correspond to a point on the cube.

Algorithm steps

  1. Turn the picture into rgb histogram, and you can imagine a color block in the space, which is R-axis, G-axis and B-axis respectively.
  2. Find the longest axis and sort with the longest axis.
  3. Divide the color block into two according to the sorted results.
  4. Continue the two color blocks obtained in step 3 to step 234 until the number of color blocks reaches the number k of extracted colors.
  5. Get k color blocks, and count the pixel RGB mean value of each color block, that is, the final extracted color result.


# k is the extracted chromatic number
def medianSegmentation(imgPath, k = 256):
	# Read picture	
	img = cv2.imdecode(np.fromfile(imgPath, dtype=np.uint8), -1)
	img = np.array(img)
	# 3D to 2D
	img = img.reshape((-1, img.shape[2]))
	# queue
	que = Queue()
	# Join the team
	# Do not stop until the cutting quantity reaches k
	while que.qsize() < k:
		# Color block
		img = que.get()
		# Find the largest axis in rgb
		selectColor = 0
		selectDifference = 0
		for i in range(3):
			selectMinVal = np.min(img[:, i])
			selectMaxVal = np.max(img[:, i])
			if selectMaxVal - selectMinVal > selectDifference:
				selectColor = i
				selectDifference = selectMaxVal - selectMinVal
		# Sort using the largest axis
		sortedIndexs = np.lexsort((img[:, selectColor], ))
		img = img[sortedIndexs, :]
		# segmentation
		leftImg = img[:img.shape[0]//2, :]
		rightImg = img[img.shape[0]//2:, :]
		# Queue
	# Extract color
	colors = []
	while not que.empty():
		# Color block
		img = que.get()
		# Take pixel mean
		colors.append(np.mean(img, axis=0))
	# Display color block
	showColors = np.zeros(shape=(len(colors) * 20, 200, 3), dtype=np.uint8)
	for i in range(k):
		showColors[i*20: i*20+20, :, :] = colors[i]
	cv2.imshow("", showColors)


  1. Because the sorting algorithm is carried out after each color block segmentation, the performance is very poor. Therefore, RGB can be sorted once and optimized in some ways.
  2. In terms of memory, repetition opens up space; It can be integrated and optimized in combination with the first point.
  3. Some articles say that VBox may have a large volume but only contain a small number of pixels under some conditions. The solution is that each time you perform segmentation, you do not segment all the VBox obtained from the last segmentation, but sort through a priority queue. At the beginning, this queue takes VBox only the number of pixels contained in VBox as the priority. When the number of segmentation times becomes more, you take volume * the number of pixels contained in VBox as the priority. VBox is the color block described in this article.


Its time complexity and space complexity have great advantages, and the fidelity is also very high. Excerpt from this paragraph Image theme color extraction algorithm , the experimental results show that the time complexity is large.

Algorithm principle

The description of the algorithm is complex. Its principle is to convert the RGB of color into binary. The binary form is xxxx xxxx, and each column of the 8-bit binary form of RGB is bonded. For example:

R: 0100 1010
G: 0101 0100
B: 0010 0101

The bonded list is:

level0: 000		corresponding index Is 0, Namely in-1 Layer node children[0]
level1: 110		corresponding index For 6, That is, at the node of layer 0 children[0]->children[6]
level2: 001		corresponding index Is 1, That is, at the node of layer 0 children[0]->children[6]->children[1]
level3: 010		corresponding index For 2, That is, at the node of layer 0 children[0]->children[6]->children[1]->children[2]

level4: 100		corresponding index For 4, That is, at the node of layer 0 children[0]->children[6]->children[1]->children[2]->children[4]
level5: 011		corresponding index For 3, That is, at the node of layer 0 children[0]->children[6]->children[1]->children[2]->children[4]->children[3]
level6: 100		corresponding index For 4, That is, at the node of layer 0 children[0]->children[6]->children[1]->children[2]->children[4]->children[3]->children[4]
level7: 001		corresponding index Is 1, That is, at the node of layer 0 children[0]->children[6]->children[1]->children[2]->children[4]->children[3]->children[4]->children[1]

For details, please refer to the following two articles: Image theme color extraction algorithm,Octree color quantization.


# Constant class
class Const(object):

# Color class
class Color(object):

	# Color constructor
	def __init__(self, r, g, b):
		self.r = r
		self.g = g
		self.b = b

	# Color addition
	def add(self, color):
		self.r += color.r
		self.g += color.g
		self.b += color.b

	# Color division
	def div(self, k):
		if k == 0:
			raise Exception("error color div zero.")
		return Color(self.r // k, self.g // k, self.b // k)

	# According to the octree principle, we can get which children are
	def getIndex(self, level):
		r = "{0:08b}".format(self.r)[level]
		g = "{0:08b}".format(self.g)[level]
		b = "{0:08b}".format(self.b)[level]
		return int(''.join([r, g, b]), 2)

	def __str__(self):
		return "Color({0}, {1}, {2})".format(self.r, self.g, self.b)

	def __repr__(self):
		return str(self)

# Octree node class
class Node(object):

	# Node constructor
	def __init__(self, level, parent):
		self.color = Color(0, 0, 0)						# Node color
		self.level = level  							# level to which the node belongs
		self.children = [None for i in range(8)]		# children owned by node
		self.pixedCount = 0  							# Number of the same color
		if level < Const.MAX_LEVEL - 1:					# If the node level is 7, it will not enter the octree levels linked list
			parent.addLevelNode(level + 1, self)		# Since the level of the root node is - 1, the + 1 operation is performed

	# Recursively create a Color path to the leaf node
	def addColor(self, color, level, parent):
		if level < Const.MAX_LEVEL:										# level less than 8
			index = color.getIndex(level)								# The level is from 0 to 7, which is the index obtained by obtaining the rgb code of each column
			if self.children[index] is None:							# The corresponding child has not been created yet
				self.children[index] = Node(level, parent)				# Create a child node and place it in the index position
			self.children[index].addColor(color, level + 1, parent)		# Recursive level to the next level
		else:															# Level reaches the 8th level, which is the leaf node
			self.color.add(color)										# Layer 8 does not need to create a Node, but directly accumulates the color of layer 7
			self.pixedCount += 1  										# Accumulate the number of color s on layer 7

	# Get leaf nodes including itself and children
	def leafNodes(self):
		leafNodes = []											# Record leaf node
		if self.isLeaf():
			leafNodes.append(self)								# append leaf node
			for node in self.children:							# Traversal subtree
				if not node is None:
					leafNodes = leafNodes + node.leafNodes()	# Get the leaf node of the subtree
		return leafNodes										# Return leaf node

	# Is it a leaf node
	def isLeaf(self):
		return self.pixedCount > 0 		# If the PixedCound of the Node is greater than 0, it indicates that it is a leaf Node

	# Merge child nodes (this operation needs to be called externally from level=7, that is, from the leaf node to the root node; in this process, the leaf (child) node of node A is merged, and node A becomes A leaf node)
	def reduce(self):
		reduceCount = 0  								# Number of merged leaf nodes
		for node in self.children:						# Traversal child
			if not node is None:
				self.color.add(node.color)				# Adds leaf node color values to the parent node
				self.pixedCount += node.pixedCount		# Add the number of leaf node pixels to the parent node
				reduceCount += 1 						# Merge count
		self.children = [None for i in range(8)]  		# Abandon the child
		return reduceCount - 1  						# Since its own Node becomes a leaf Node, + 1

	# The Color value on the Node
	def normalize(self):
		return self.color.div(self.pixedCount)			# Mean Color

# Octree algorithm
class Octree(object):

	# Octree constructor
	def __init__(self):
		self.levels = [[] for i in range(Const.MAX_LEVEL)] 		# Build a levels linked list for subsequent extraction of color theme colors
		self.root = Node(-1, self)								# root node

	# Add Node to the levels linked list
	def addLevelNode(self, level, node):

	# Add color to octree
	def addColor(self, color):
		self.root.addColor(color, 0, self)

	# Extract color
	def extractColor(self, k = 256):
		leafCount = len(self.root.leafNodes())		# Get the number of octree leaf nodes
		for i in range(Const.MAX_LEVEL, 0, -1):		# Traverse from level 7 and merge leaf nodes
			level = i - 1   						# Since i starts at 8, - 1
			if leafCount <= k:						# If the number of leaf nodes is less than the extracted number k, end
			if not self.levels[level] is None:		# The linked list is not empty. Traverse the parent layer of the leaf node, because the color of the parent node is to be counted here
				for node in self.levels[level]:		# Traverse the node node in the linked list
					leafCount -= node.reduce()		# Count the color of node
					if leafCount <= k:				# If the number of leaf nodes is less than the extracted number k, end
			self.levels[level] = []					# level null
		# Extraction color
		colors = []
		# Get all leaf nodes of the merged octree
		for leafNode in self.root.leafNodes():
			if leafNode.isLeaf() and len(colors) <= k:
				# Get the average value of Node color
		# Returns the extracted color
		return colors

def octreeColor(imgPath):
	# Read picture
	img = cv2.imdecode(np.fromfile(imgPath, dtype=np.uint8), -1)
	img = np.array(img)
	width, height, channel = img.shape
	# Create octree
	octree = Octree()
	# Add color
	for i in range(width):
		for j in range(height):
			octree.addColor(Color(img[i, j, 0], img[i, j, 1], img[i, j, 2]))
	# Extract color
	k = 16
	colors = octree.extractColor(k)
	# Display color block
	showColors = np.zeros(shape=(len(colors) * 20, 200, 3), dtype=np.uint8)
	for i in range(len(colors)):
		c = [colors[i].r, colors[i].g, colors[i].b]
		showColors[i*20: i*20+20, :, :] = np.array(c)
	cv2.imshow("", showColors)


After testing: about 20 ~ 30 seconds on your PC platform.



Color modeling


Project link

Theme color extraction project

Posted by Mark.P.W on Fri, 05 Nov 2021 14:47:49 -0700