marker_separation/py/common.py

import cv2
import numpy as np
from math import sqrt

def convert_color (color, color_space_a, color_space_b) :
	pixel = np.zeros([1, 1, 3], dtype=np.uint8)

	if color_space_a == 'RGB' :
		pixel = cv2.cvtColor(pixel, cv2.COLOR_BGR2RGB)
	elif color_space_a == 'LAB' :
		pixel = cv2.cvtColor(pixel, cv2.COLOR_BGR2LAB)
	elif color_space_a == 'HSV' :
		pixel = cv2.cvtColor(pixel, cv2.COLOR_BGR2HSV)
	
	#default is BGR
	pixel[:] = color

	if color_space_a == 'RGB' and color_space_b == 'BGR' :
		b = cv2.COLOR_RGB2BGR
	elif color_space_a == 'BGR' and color_space_b == 'RGB' :
		b = cv2.COLOR_BGR2RGB
	elif color_space_a == 'RGB' and color_space_b == 'LAB' :
		b = cv2.COLOR_RGB2LAB
	elif color_space_a == 'LAB' and color_space_b == 'RGB' :
		b = cv2.COLOR_LAB2RGB
	elif color_space_a == 'BGR' and color_space_b == 'LAB' :
		b = cv2.COLOR_BGR2LAB
	elif color_space_a == 'LAB' and color_space_b == 'BGR' :
		b = cv2.COLOR_LAB2BGR
	elif color_space_a == 'HSV' and color_space_b == 'LAB' :
		b = cv2.COLOR_HSV2LAB
	elif color_space_a == 'LAB' and color_space_b == 'HSV' :
		b = cv2.COLOR_LAB2HSV
	elif color_space_a == 'RGB' and color_space_b == 'HSV' :
		b = cv2.COLOR_RGB2HSV
	elif color_space_a == 'HSV' and color_space_b == 'RGB' :
		b = cv2.COLOR_HSV2RGB
	elif color_space_a == 'BGR' and color_space_b == 'HSV' :
		b = cv2.COLOR_BGR2HSV
	elif color_space_a == 'HSV' and color_space_b == 'BGR' :
		b = cv2.COLOR_HSV2BGR
	elif color_space_a == 'RGB' and color_space_b == 'RGB' :
		b = None
	elif color_space_a == 'BGR' and color_space_b == 'BGR' :
		b = None
	elif color_space_a == 'LAB' and color_space_b == 'LAB' :
		b = None
	elif color_space_a == 'HSV' and color_space_b == 'HSV' :
		b = None

	if b is not None :
		cvt = cv2.cvtColor(pixel, b)
	else :
		cvt = pixel

	return cvt[0, 0]

def closest_color (colors, color):
	colors = np.array(colors)
	color = np.array(color)
	distances = np.sqrt(np.sum((colors - color) ** 2, axis=1))
	index_of_smallest = np.where(distances == np.amin(distances))
	smallest_distance = colors[index_of_smallest]
	return smallest_distance[0]

# Works for RGB, BGR, LAB and HSV(?)
def closest_color_pythagorean (colors, color) :
	mDist = float('inf')
	mIdx = -1
	for idx, comp in enumerate(colors) :
		dist = pythagorean_distance(comp[0], comp[1], comp[2], color[0], color[1], color[2])
		if dist < mDist :
			mDist = dist
			mIds = idx
	return color[mIdx], mDist

def closest_color_weighted_pythagorean (colors, color, space) :
	mDist = float('inf')
	mIdx = -1
	for idx, comp in enumerate(colors) :
		if space == 'BGR' :
			dist = weighed_pythagorean_distance(comp[2], comp[1], comp[0], color[2], color[1], color[1])
		elif space == 'RGB' :
			dist = weighed_pythagorean_distance(comp[0], comp[1], comp[2], color[0], color[1], color[2])
		else :
			raise Exception(f'closest_color_weighted_pythagorean does not support color space {space}')
			break

		if dist < mDist :
			mDist = dist
			mIds = idx
	return color[mIdx], mDist

def create_colored_image (width, height, bgr_color):
	image = np.zeros((height, width, 3), np.uint8)
	image[:] = bgr_color
	return image

def remove_from_list (l, item) :
	new_array = []
	for i in l :
		if not list_match(i, item) :
			new_array.append(i)
	return new_array

def list_match (a, b) :
	for i in range(len(a)) :
		if a[i] != b[i] :
			return False
	return True

def pythagorean_distance (r1, g1, b1, r2, g2, b2) :
	return sqrt(pow(r1-r2, 2) + pow(g1-g2, 2) + pow(b1-b2, 2))

def weighted_pythagorean_distance (r1, g1, b1, r2, g2, b2) :
	R = 0.30
	G = 0.59
	B = 0.11
	return sqrt(pow((r1-r2) * R, 2) + pow((g1-g2) * G, 2) + pow((b1-b2) * B, 2))