marker_separation/py/pallete.py

import argparse
import cv2
import numpy as np
from math import floor
from pallete_schema import PalleteSchema

from os.path import isfile, realpath, basename, dirname, splitext, join

parser = argparse.ArgumentParser(description='Separate an image into most similar colors specified')

parser.add_argument('input', type=str, help='Input image to extract the pallete from')

class Pallete :
	input = ''
	output = ''
	image = None
	selection = []
	colors = []
	i = 0

	def __init__ (self, args) :
		if isfile(args.input) :
			self.input = realpath(args.input)
		else :
			print(f'File {self.input} does not exist')
			exit(1)

		self.set_output()
		self.load()
		self.display_ui()
		self.write()

	def set_output (self) :
		dir = dirname(self.input)
		stem = splitext(basename(self.input))[0]
		name = basename(self.input)
		print(f'Using image {name}')
		self.output = join(dir, f'{stem}.json')
		print(f'Writing to {stem}.json')

	def load (self) :
		self.image = cv2.imread(self.input)

	def display_ui (self) :
		cv2.imshow('pallete', self.image)
		cv2.namedWindow('pallete')
		cv2.setMouseCallback('pallete', self.on_click)
		cv2.waitKey(0)
		cv2.destroyAllWindows()

	def on_click (self, event, x, y, flags, param) :
		if event == cv2.EVENT_LBUTTONDOWN:
			#print(f'{x},{y}')
			if len(self.selection) == 0 :
				self.selection.append([x, y])
			elif len(self.selection) == 1 :
				self.selection.append([x, y])
				self.process_color(self.selection)
				self.selection = []

	def process_color (self, selection) :
		top_left = selection[0]
		bottom_right = selection[1]
		color_select = self.image[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
		avg_color_per_row = np.average(color_select, axis=0)
		avg_color = np.average(avg_color_per_row, axis=0)
		int_color = []
		for i in avg_color :
			int_color.append( int( round(i) ) )
		self.colors.append({
			'name' : str(self.i),
			'color' : int_color,
			'space' : 'BGR'
		})
		print(f'Added {self.i} = {int_color}')
		self.i += 1

	def write (self) :
		PalleteSchema(self.output, self.colors)


if __name__ == "__main__" :
	args = parser.parse_args()
	Pallete(args)