import PIL.Image as Image
import PIL.ImageDraw as ImageDraw
import colorsys
import random

import scipy.misc
import scipy.fftpack
import scipy

import math

import numpy

import pyflann

def get_sample(image, x, y):
	block_size = 8
	halfsize = block_size // 2
	
	if x < halfsize:
		x = halfsize
	if y < halfsize:
		y = halfsize
	if x >= image.size[0] - halfsize:
		x = image.size[0] - halfsize - 1
	if y >= image.size[1] - halfsize:
		y = image.size[1] - halfsize - 1
	
	tile = image.crop((x - halfsize, y - halfsize, x + halfsize, y + halfsize))
	sample = numpy.array(tile).ravel()
	sample -= sample.sum() / sample.size
	
	return sample

wood = Image.open('wood.png').convert('RGB')

input_count = len(get_sample(wood, 0, 0))

print 'Sample size:', input_count

flann = pyflann.FLANN()

class Status:
	def __init__(self):
		self.status = 0
	
	def __call__(self, now, total):
		newstatus = int(10 * now / total) * 10
		if newstatus != self.status:
			self.status = newstatus
			print str(newstatus) + "%",

status = Status()
			
print "Collecting training samples..",
samples = []
outputs = []
for y in range(0, wood.size[1], 4):
	status(y, wood.size[1])
	for x in range(0, wood.size[0], 4):
		sample = get_sample(wood, x, y)
		
		sample = tuple(sample)
		samples.append(sample)
		
print "Got", len(samples), "samples"

print "Wood:", len([value for value in outputs if value == 1]),
print "Not wood:", len([value for value in outputs if value == -1])

print "Building index..",		
flann.build_index(numpy.array(samples), algorithm = 'kdtree')
print "done"

print "Collecting test samples..",
test = Image.open('wms_nuuksio2.png').convert('RGB')
output = Image.new('L', test.size)
draw = ImageDraw.Draw(output)
step = 4
for y in range(0, test.size[1], step):
	status(y, test.size[1])
	
	samples = []
	for x in range(0, test.size[0], step):
		sample = get_sample(test, x, y)
		samples.append(sample)
	
	result, dists = flann.nn_index(numpy.array(samples), 3)
	
	i = 0
	for x in range(0, test.size[0], step):
		value = math.sqrt(dists[i][0])
		draw.rectangle((x,y,x + step, y + step), fill = int(255 - value / 4.))
		i += 1
	
print "done"
		
output.save('output.png')