Skip to content

Instantly share code, notes, and snippets.

@hirocarma

hirocarma/color_cluster.py

Created February 28, 2021 06:01
Show Gist options
  • Save hirocarma/ed440718f0f612af7edbf4855113e033 to your computer and use it in GitHub Desktop.
Save hirocarma/ed440718f0f612af7edbf4855113e033 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
color_cluster.py
#!/usr/bin/env python
import sys
import os
import shutil
import math
import matplotlib.pyplot as plt
import cv2
import numpy as np
def rotate(img, angle):
h, w = img.shape[:2]
size = (w, h)
angle_rad = angle/180.0*np.pi
w_rot = int(np.round(h*np.absolute(np.sin(angle_rad))\
+w*np.absolute(np.cos(angle_rad))))
h_rot = int(np.round(h*np.absolute(np.cos(angle_rad))\
+w*np.absolute(np.sin(angle_rad))))
size_rot = (w_rot, h_rot)
center = (w/2, h/2)
rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
affine_matrix = rotation_matrix.copy()
affine_matrix[0][2] = affine_matrix[0][2] -w/2 + w_rot/2
affine_matrix[1][2] = affine_matrix[1][2] -h/2 + h_rot/2
img_rot = cv2.warpAffine(img, affine_matrix, size_rot, flags=cv2.INTER_CUBIC)
return img_rot
def scale_to_resolation(img, resolation):
h, w = img.shape[:2]
scale = (resolation / (w * h)) ** 0.5
return cv2.resize(img, dsize=None, fx=scale, fy=scale)
def show_imgs_bfmatch(img1,img2):
akaze = cv2.AKAZE_create()
kp1, des1 = akaze.detectAndCompute(img1, None)
kp2, des2 = akaze.detectAndCompute(img2, None)
bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
matches = bf.match(des1, des2)
matches = sorted(matches, key = lambda x:x.distance)
img3 = cv2.drawMatches(img1, kp1, img2, kp2, matches[:10], None, flags=2)
cv2.imshow('match', cv2.cvtColor(img3, cv2.COLOR_BGR2RGB))
def bfmatch(target_img, comp_img, DEBUG):
target_img_size = (target_img.shape[1],target_img.shape[0] )
comp_img = scale_to_resolation(comp_img, \
target_img.shape[1] * target_img.shape[0])
comp_img = cv2.GaussianBlur(comp_img,(5,5),0)
comp_img = cv2.medianBlur(comp_img, 3)
comp_img = cv2.bilateralFilter(comp_img, 9, 75, 75)
#comp_img = cv2.fastNlMeansDenoising(comp_img, h=30)
if DEBUG == 1: show_imgs_bfmatch(target_img,comp_img)
bf = cv2.BFMatcher(cv2.NORM_HAMMING)
detector = cv2.AKAZE_create()
(target_kp, target_des) = detector.detectAndCompute(target_img, None)
if DEBUG == 1:
cv2.imshow("target_image",target_img)
cv2.waitKey(0)
cv2.imshow("comp_image",comp_img)
cv2.waitKey(0)
try:
(comparing_kp, comparing_des) = detector.detectAndCompute(comp_img, None)
matches = bf.match(target_des, comparing_des)
dist = [m.distance for m in matches]
ret = sum(dist) / len(dist)
except cv2.error:
ret = 100000
return ret
def hue_check(img, upper, lower):
div = 30 ; per = div * 0.5
i = 0
height, width, chan = img.shape
hClp = height / div
wClp = width / div
for x in range(0,div):
for y in range(0,div):
Clp = img[int(y * hClp):int((y + 1) * hClp), \
int(x * wClp):int((x + 1) * wClp)]
h_c, s_c, v_c = Clp[:,:,0], Clp[:,:,1], Clp[:,:,2]
med_h_c = np.mean(Clp[:,:,0])
med_s_c = np.mean(Clp[:,:,1])
med_v_c = np.mean(Clp[:,:,2])
if med_h_c < upper and med_h_c > lower:
i = i + 1
if i > per:
return True
return False
def calc_degree(a,b,c):
vec_a = a - b
vec_c = c - b
#cos
length_vec_a = np.linalg.norm(vec_a)
length_vec_c = np.linalg.norm(vec_c)
inner_product = np.inner(vec_a, vec_c)
cos = inner_product / (length_vec_a * length_vec_c)
#radian
rad = np.arccos(cos)
#degree
return np.rad2deg(rad)
def color_cluster(src_img, target_img, DEBUG):
hsv_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2HSV_FULL)
h, s, v = hsv_img[:,:,0], hsv_img[:,:,1], hsv_img[:,:,2]
med_h = np.mean(hsv_img[:,:,0])
med_s = np.mean(hsv_img[:,:,1])
med_v = np.mean(hsv_img[:,:,2])
if DEBUG == 1:
print ('med_hsv:', med_h, med_s, med_v)
div = 180
hight, width, chan = hsv_img.shape
hClp = hight / div
wClp = width / div
FLG = 0
for x in range(0,div):
for y in range(0,div):
hsv_Clp = \
hsv_img[int(y * hClp):int((y + 1) * hClp), \
int(x * wClp):int((x + 1) * wClp)]
h_c, s_c, v_c = hsv_Clp[:,:,0], hsv_Clp[:,:,1], hsv_Clp[:,:,2]
med_h_c = np.median(hsv_Clp[:,:,0])
med_s_c = np.median(hsv_Clp[:,:,1])
med_v_c = np.median(hsv_Clp[:,:,2])
if (med_h_c < 8 or med_h_c > 230) and \
med_s_c < 148 and med_s_c > 128: #med_v_c > 50:
FLG = FLG + 1 ; break
if FLG == 2: break
if DEBUG == 1:
print('med_h_c med_s_c med_v_c ', med_h_c, med_s_c, med_v_c)
print('x y ', str(x * wClp) , str(y * hClp))
check_img = src_img.copy()
cv2.rectangle(check_img, (int(x * wClp), int(y * hClp)), \
(int(x * wClp + wClp) , int(y * hClp + hClp)), (0,0,255), 1)
cv2.imshow("check", check_img)
cv2.waitKey(0)
if FLG > 1:
hl_save, hu_save = 0, 0
hl = med_h_c - 16 ; hl = min(255, max(hl, 0))
hu = med_h_c + 10
if hu > 255:
hl_save = hl - 12; hu_save = 255
hl = 0 ; hu = hu -255 + 8
hu = min(255, max(hu, 0))
sl = med_s_c - 58 ; sl = min(255, max(sl, 0))
su = med_s_c + 28 ; su = min(255, max(su, 0))
vl = med_v_c - 20 ; vl = min(255, max(vl, 0))
vu = med_v_c + 48 ; vu = min(255, max(vu, 0))
else:
return src_img
if DEBUG == 1:
print('hl sl vl ', hl , sl, vl )
print('hu su vu ', hu , su, vu )
if hl_save > 0 : print('hl_save hu_save', hl_save , hu_save )
lower = np.array([hl,sl,vl], dtype=np.uint8)
upper = np.array([hu,su,vu], dtype=np.uint8)
mask = cv2.inRange(hsv_img, lower, upper)
if hl_save > 0:
lower = np.array([hl_save,sl,vl], dtype=np.uint8)
upper = np.array([hu_save,su,vu], dtype=np.uint8)
mask = cv2.inRange(hsv_img, lower, upper)
kernel = np.ones((5,5),np.uint8)
mask = cv2.dilate(mask,kernel,iterations = 4)
mask = cv2.erode(mask,kernel,iterations = 4)
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
mask = cv2.dilate(mask,kernel,iterations = 2)
mask = cv2.medianBlur(mask, 3)
if DEBUG == 1:
cv2.imshow("mask_image", mask)
cv2.waitKey(0)
contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, \
cv2.CHAIN_APPROX_SIMPLE)
filtered_contour = []
result_img = src_img.copy()
for indx in range(len(contours)):
cnt = contours[indx]
if DEBUG == 1 : print('[ID]',indx)
if len(cnt) > 20: #160
area = cv2.contourArea(cnt)
if area < len(cnt):
continue
if area < 2056 or area > 410000:
if DEBUG == 1 : print('Exit:area',area)
continue
rect = cv2.minAreaRect(cnt)
center, size, angle = rect
size_x, size_y = size #clean
if size_x < 16 or size_y < 11:
if DEBUG == 1 : print('Exit:size_x size_y',size_x,size_y)
continue
rec_area = size_x * size_y
area_ratio = area / rec_area
if area_ratio < 0.12:
if DEBUG == 1 : print('Exit:area_ratio',area_ratio)
continue
len_ratio = len(cnt) / ((size_x + size_y) * 2)
if len_ratio < 0.26:
if DEBUG == 1 : print('Exit:len_ratio',len_ratio)
continue
np_contour = \
np.array(cnt).reshape(len(cnt),2)
left_x = min(np_contour[:,0])
right_x = max(np_contour[:,0])
top_y = min(np_contour[:,1])
bottom_y = max(np_contour[:,1])
v_ratio = (right_x - left_x) / (bottom_y - top_y)
if v_ratio < 0.6 or v_ratio > 5:
if DEBUG == 1 : print('Exit:v_ratio',v_ratio)
continue
v_rec_w = right_x - left_x ; v_rec_h = bottom_y - top_y
v_rec_area = v_rec_w * v_rec_h
v_area_ratio = area / v_rec_area
if v_area_ratio < 0.12 or v_rec_w < 118 or v_rec_h < 34:
if DEBUG == 1 : print('Exit:v_area_ratio',v_area_ratio,v_rec_w,v_rec_h)
continue
#
mask_rec_img = mask[top_y:bottom_y,left_x:right_x]
if DEBUG == 1:
cv2.imshow("mask_rec_img1", mask_rec_img)
cv2.waitKey(0)
if abs(angle) > 45:
rec_angle = 90 - abs(angle)
else:
rec_angle = angle
mask_rec_img = rotate(mask_rec_img, rec_angle)
if DEBUG == 1:
cv2.imshow("mask_rec_img2", mask_rec_img)
cv2.waitKey(0)
mask_h, mask_w = mask_rec_img.shape
if size_y > size_x:
v_x = size_y ; v_y = size_x
else:
v_x = size_x ; v_y = size_y
m_top_y = int(mask_h/2 - v_y/2)
m_bottom_y = int(mask_h/2 + v_y/2)
m_left_x = int(mask_w/2 - v_x/2)
m_right_x = int(mask_w/2 + v_x/2)
mask_rec_img = mask_rec_img[m_top_y:m_bottom_y,m_left_x:m_right_x]
if DEBUG == 1:
cv2.imshow("mask_rec_img3", mask_rec_img)
cv2.waitKey(0)
bf_len = bfmatch(target_img, mask_rec_img, DEBUG)
if bf_len > 64:
if DEBUG == 1 : print('Exit:bf_len',bf_len)
continue
#
rec_img = result_img[top_y:bottom_y,left_x:right_x]
if not hue_check(rec_img, 60, 20): #not general
if DEBUG == 1 : print('Exit:hue_check')
continue
box = cv2.boxPoints(rect)
p1_x, p1_y = box[1]
p0_x, p0_y = box[0]
p_a = np.array([p1_x,p0_y])
p_b = np.array([p1_x,p1_y])
p_c = np.array([p0_x,p0_y])
degree = calc_degree(p_a,p_b,p_c)
if degree < 62 and degree > 28:
if DEBUG == 1 : print('Exit:degree', degree)
continue
if size_x < size_y or math.isnan(degree): #depend degree
rec_ratio = size_y / size_x
else:rec_ratio = size_x / size_y
if rec_ratio < 0.8 or rec_ratio > 4:
if DEBUG == 1 : print('Exit:rec_ratio',rec_ratio)
continue
box = np.int0(box)
result_img = cv2.drawContours(result_img,[box],0,(0,0,255),2)
filtered_contour.append(cnt)
if DEBUG == 1:
print('len:', len(cnt))
print('center', center) ; print('size', size)
print('angle', angle) ; print('area', area)
print('rec_area', rec_area) ; print('area_ratio', area_ratio)
print('rec_ratio', rec_ratio) ; print('len_ratio',len_ratio)
print('v_ratio', v_ratio) ; print('v_area_ratio',v_area_ratio)
print('degree', degree) ; print('top_y',top_y)
print('bf_len',bf_len)
result_img = cv2.line(result_img, (p1_x,p1_y), \
(p0_x,p0_y), (255, 255, 255),1)
result_img = cv2.line(result_img, (p1_x,p1_y), \
(p1_x,p0_y), (255, 255, 255),1)
cv2.putText(result_img, '['+ str(round(degree,3)) + ']', \
(int(p1_x) - 60, int(p1_y) + 20 ), \
cv2.FONT_HERSHEY_PLAIN, 0.8, (255, 255, 255), \
1, cv2.LINE_AA)
cv2.putText(result_img, '['+ str(round(p1_x,1)) + 'x' + \
str(round(p1_y,1))+ ']', \
(int(p1_x) , int(p1_y)), \
cv2.FONT_HERSHEY_PLAIN, 0.8, (255, 255, 255), \
1, cv2.LINE_AA)
cv2.putText(result_img, '['+ str(round(p0_x,1)) + 'x' + \
str(round(p0_y,1))+ ']', \
(int(p0_x) , int(p0_y)), \
cv2.FONT_HERSHEY_PLAIN, 0.8, (255, 255, 255), \
1, cv2.LINE_AA)
center_x, center_y = center
cv2.rectangle(result_img, (left_x,top_y), \
(right_x, bottom_y), (255,0,0), 2)
cv2.putText(result_img, '['+ str(indx) +']', \
(int(center_x), int(center_y)), \
cv2.FONT_HERSHEY_PLAIN, 0.8, (255, 255, 255), \
1, cv2.LINE_AA)
pass
result_img = cv2.drawContours(result_img, filtered_contour, -1, (0,255,0), 1)
if DEBUG == 1:
cv2.imshow("result",result_img)
cv2.waitKey(0)
return result_img
if __name__ == "__main__":
TARGET_FILE = '/home/hiro/target-megane.jpg'
target_img = cv2.imread(TARGET_FILE)
IMG_FILE = sys.argv[1]
if (len(sys.argv)) == 3:
OUT_DIR = sys.argv[2]
if os.path.isfile(IMG_FILE):
DEBUG = 1
src_img = cv2.imread(IMG_FILE)
result_img = color_cluster(src_img, target_img, DEBUG)
sys.exit()
if os.path.isdir(OUT_DIR):
shutil.rmtree(OUT_DIR)
os.makedirs(OUT_DIR)
if os.path.isdir(IMG_FILE):
DEBUG = 0
IMG_DIR = IMG_FILE
files = os.listdir(IMG_DIR)
files = sorted(files)
for file in files:
src_img_path = IMG_DIR + file
src_img = cv2.imread(src_img_path)
result_img = color_cluster(src_img, target_img, DEBUG)
fname = file.split('.')[0]
cv2.imwrite(OUT_DIR + 'megane-' + fname + '.jpg', result_img)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment