85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
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)
|