import logging as log
import numpy as np
import scipy.cluster
import cv2 as cv
def kmeans(img):
arr=np.reshape(img,(-1,3)).astype(np.float)
wood=[193,165,116]
(centers,distortion)=scipy.cluster.vq.kmeans(arr,3)
log.debug("k-means centers: %s",centers)
(black,empty,white)=sorted(centers,key=sum)
if np.linalg.norm(black)>np.linalg.norm(black-wood):
black=None
if np.linalg.norm(white-[255,255,255])>np.linalg.norm(white-wood):
white=None
log.debug("black, white: %s, %s",black,white)
return (black,white,centers)
class QuantizedImage:
BLACK=0
WHITE=1
EMPTY=2
def __init__(self,img):
self.img=self._quantize(img)
self._mask()
def transform(self,matrix):
(h,w)=self.img.shape[:2]
self.img=cv.warpPerspective(self.img,matrix,(w,h))
self.maskB=cv.warpPerspective(self.maskB,matrix,(w,h))
self.maskW=cv.warpPerspective(self.maskW,matrix,(w,h))
def get(self,x,y):
if self.maskB[y,x]: return self.BLACK
elif self.maskW[y,x]: return self.WHITE
else: return self.EMPTY
def _quantize(self,img):
(self._black,self._white,colors)=self._sampleColors(img)
origShape=img.shape
data=np.reshape(img,(-1,3))
(keys,dists)=scipy.cluster.vq.vq(data,colors)
pixels=np.array([colors[k] for k in keys],dtype=np.uint8).reshape(origShape)
return pixels
def _sampleColors(self,rect):
(h,w)=rect.shape[:2]
minirect=rect[h//4:3*h//4, w//4:3*w//4]
return kmeans(minirect)
def _mask(self):
unit=np.array([1,1,1],dtype=np.uint8)
if self._black is not None:
self.maskB=cv.inRange(self.img,self._black-unit,self._black+unit)
else:
self.maskB=np.zeros(self.img.shape[:2],dtype=np.uint8)
if self._white is not None:
self.maskW=cv.inRange(self.img,self._white-unit,self._white+unit)
else:
self.maskW=np.zeros(self.img.shape[:2],dtype=np.uint8)