#!/usr/bin/python

# segment.py
# Part of the "penpal" project
# (Finds Gary Blehm's Identical Penmen)
# http://web.mit.edu/wdaher/www/penpal/ 
#
# Waseem Daher
# Jan. 21, 2006

import Image, ImageChops
import os, sys
import time

#infile = '4x4-identical-bw.jpg'
infile = sys.argv[1]
datafile = 'penmen-data.txt'

# What value do we consider our
# threshold for 'black'?
BLACK_THRESHOLD = 128
# Throw away objects whose dimensions
# are < TOO_SMALL times the average 
TOO_SMALL = 1
# Throw away objects w/o enough pixels
TOO_FEW_PIXELS = 1
# When selecting a bounding box, make it
# BOUNDING_BOX times bigger than the average
# penman
BOUNDING_BOX_W = 2.5
BOUNDING_BOX_H = 3
# What percentage of the avg. to move
# up the y-val of the centroid?
# (because we want to capture the head)
HEIGHT_SHIFT = 0.3
# Magic numbers are a hack, but hey
BIG_NUMBER = 99999999999

class Penman:
    def __init__(self):
        self.points = []
        self.centroid = None
        self.dims = None
    def addPoint(self, point):
        self.points.append(point)
    def getCentroid(self):
        if self.centroid is None:
            self._getCentroidAndDims()
        return self.centroid
    def getDimensions(self):
        if self.dims is None:
            self._getCentroidAndDims()
        return self.dims
    def getNumPixels(self):
        return len(self.points)
    def _getCentroidAndDims(self):
        sumx = sumy = 0
        # Magic numbers are sort of a hack,
        # but perhaps this is faster than
        # doing an or comparison every time
        minx = miny =  BIG_NUMBER
        maxx = maxy = -BIG_NUMBER
        for point in self.points:
            (x,y) = point
            if x < minx:
                minx = x
            if x > maxx:
                maxx = x
            if y < miny:
                miny = y
            if y > maxy:
                maxy = y
            sumx += x
            sumy += y
        avgx = sumx / len(self.points)
        avgy = sumy / len(self.points)
        self.centroid = (avgx,avgy)
        self.dims = (maxx-minx, maxy-miny)

def getNeighbors(x,y,rows,cols):
    """
    Returns the set of points
    that are the neighbors of the
    supplied point
    """
    neighbors = []
    if x > 0 and y > 0:         
        neighbors.append( (x-1,y-1) )
    if x > 0:                   
        neighbors.append( (x-1,y) )
    if x > 0 and y<rows-1:      
        neighbors.append( (x-1,y+1) )

    if y > 0:                   
        neighbors.append( (x,y-1) )
    if y < rows-1:              
        neighbors.append( (x,y+1) )

    if x < cols-1 and y > 0:    
        neighbors.append( (x+1,y-1) )
    if x < cols-1 and y > 0:    
        neighbors.append( (x+1,y) )
    if x < cols-1 and y<rows-1: 
        neighbors.append( (x+1,y+1) )

    return neighbors

def findPenman(x,y,im):
    """
    Given a starting black pixel, find the
    maximally-sized connected region of
    pixels. This is a destructive read -
    the pixels get set to 255 after we 
    go over them.
    """
    (cols,rows) = im.size
    me = Penman()
    stack = [(x,y)]
    while len(stack) > 0:
        (a,b) = stack.pop()
        # Blank the pixel so we don't try to
        # process this one again
        im.putpixel((a,b), 255)
        me.addPoint((a,b))
        # Check all the neighbors
        # and add them if they are
        # black.
        neighbors = getNeighbors(a,b,rows,cols)
        for (i,j) in neighbors:
            try:
                if im.getpixel((i,j))<BLACK_THRESHOLD:
                    stack.append((i,j))
            except:
                print "Size is:", im.size
                print "Tried to get:", i, j
    return me

im = Image.open(infile)
(cols, rows) = im.size

average_width = 0
average_height = 0
average_num_pixels = 0
myPenmen = []
for x in range(cols):
    for y in range(rows):
        if im.getpixel( (x,y) ) < BLACK_THRESHOLD:
            man = findPenman(x,y,im)
            w,h = man.getDimensions()
            average_width += w
            average_height += h
            average_num_pixels+=man.getNumPixels()
            myPenmen.append(man)

#--- Throw away "penmen" that don't make
#    the cut (e.g. too small)
average_width /= len(myPenmen)
average_height /= len(myPenmen)
average_num_pixels /= len(myPenmen)

min_num_pixels = average_num_pixels * TOO_FEW_PIXELS
min_width = average_width * TOO_SMALL
min_height = average_height * TOO_SMALL
bounding_box_width = average_width * BOUNDING_BOX_W
bounding_box_height = average_height * BOUNDING_BOX_H

realPenmen = []
for man in myPenmen:
    w,h = man.getDimensions()
    numPix = man.getNumPixels()
    if w > min_width and h > min_height \
        and numPix > min_num_pixels:
        realPenmen.append(man)

# Great, we have all the real penmen
# Dump this to a list.
file = open(datafile, 'w')
lines = []

lines.append('# Penman image details')
lines.append('# Date: %s' % time.asctime())
lines.append('infile = "%s"' % infile)
lines.append('min_width = %i' % min_width)
lines.append('min_height = %i' % min_height)
lines.append('average_width = %i' % average_width)
lines.append('average_height = %i' % average_height)
lines.append('bounding_box_width = %i' \
            % bounding_box_width)
lines.append('bounding_box_height = %i' \
            % bounding_box_height)
lines.append('num_penmen = %i' % len(realPenmen))
lines.append('')
lines.append('# Penman data')
lines.append('# Number, centroid_x, centroid_y')
for i in range(len(realPenmen)):
    centroid_x,centroid_y = realPenmen[i].getCentroid()
    lines.append('%i,%i,%i' % (i, centroid_x, \
        centroid_y - (HEIGHT_SHIFT * average_height)))

file.writelines([foo + '\n' for foo in lines])
file.close()

#os.system('cat %s' % datafile)