draw.py

from PIL import Image, ImageColor, ImageDraw

__colormap = (
   (255,   0,   0), (255,  31,   0), (255,  63,   0), (255,  95,   0), (255, 127,   0),
   (255, 159,   0), (255, 191,   0), (255, 223,   0), (255, 255,   0), (223, 255,   0),
   (191, 255,   0), (159, 255,   0), (127, 255,   0), ( 95, 255,   0), ( 63, 255,   0),
   ( 31, 255,   0), (  0, 255,   0), (  0, 255,  31), (  0, 255,  63), (  0, 255,  95),
   (  0, 255, 127), (  0, 255, 159), (  0, 255, 191), (  0, 255, 223), (  0, 255, 255),
   (  0, 223, 255), (  0, 191, 255), (  0, 159, 255), (  0, 127, 255), (  0,  95, 255),
   (  0,  63, 255), (  0,  31, 255), (  0,   0, 255), ( 31,   0, 255), ( 63,   0, 255),
   ( 95,   0, 255), (127,   0, 255), (159,   0, 255), (191,   0, 255), (223,   0, 255),
   (255,   0, 255), (255,   0, 223), (255,   0, 191), (255,   0, 159), (255,   0, 127),
   (255,   0,  95), (255,   0,  63), (255,   0,  31), (255, 255, 255), (  0,   0,   0)
)

def draw( points, centroids, width=0, height=0, filename="temp.bmp" ):
    _width = 0
    _height = 0
    clusters = 0

    # find width, heigth based on max x, y of points, centroids
    for p in points + centroids:
        if p.cluster > clusters: clusters = p.cluster
        if p.x > _width: _width = p.x
        if p.y > _height: _height = p.y 

    clusters += 1
    _width += 1
    _height += 1

    # use out values if user hasn't specified any dimensions
    if width==0 or height==0:
        width = _width
        height = _height

    # create the image
    image = Image.new( "RGB", (width,height), (255,255,255) )
    draw = ImageDraw.Draw( image )

    colors = [ __colormap[i] for i in range( 0, len(__colormap), int(len(__colormap)/clusters)) ]

    # draw the points
    for p in points:
        draw.ellipse( (p.x-2, p.y-2, p.x+2, p.y+2), fill=colors[p.cluster] )
        
    # draw the clusters
    for c in centroids:
        draw.ellipse( (c.x-5, c.y-5, c.x+5, c.y+5), outline=colors[c.cluster] )

    image.save( filename )