rasterize

#!/usr/bin/env python

import shapefile
import Image, ImageDraw
import sys
import netCDF4 as nc

def rasterize(filename, resolution, bg,  x_range = None, y_range = None):
    # Read in a shapefile
    data = shapefile.Reader(filename)
    # domain width and height
    if x_range is None or y_range is None:
        x_min, x_max = data.bbox[0], data.bbox[2]
        y_min, y_max = data.bbox[1], data.bbox[3]
    else:
        x_min, x_max = x_range
        y_min, y_max = y_range

    print x_min, x_max
    print y_min, y_max
    width=int(round((x_max-x_min)/resolution))
    height=int(round((y_max-y_min)/resolution))
    img = Image.new("RGB", (width, height), bg)
    draw = ImageDraw.Draw(img)


    xdist = x_max - x_min
    ydist = y_max - y_min

    xratio = width / xdist
    yratio = height / ydist

    def num2col(num):
        string = "rgb(%i,%i,%i)"%(min(255,num/256/256),(num/256)%256,num%256)
        return string
        
    def rasterize_poly(points, color):
        pixels = []
        for x,y in points:
            px = int((x - x_min) * xratio)
            py = int((y - y_min) * yratio)
            pixels.append((px,py))

        draw.polygon(pixels, fill=color)

    def rasterize_shape(shape, bg, fg):
        parts = shape.parts

        if len(parts) == 1:
            rasterize_poly(shape.points, fg)
        else:
            for j in range(0, 1) # len(parts)-1):
                if j != 0:
                    fg = bg

                start = parts[j]
                end = parts[j+1]
                rasterize_poly(shape.points[start:end], fg)

    for (num,shape) in enumerate(data.shapes()):
        rasterize_shape(shape, bg, num2col(num))

    return img

if __name__ == "__main__":
    if len(sys.argv) == 3:
        nc_file = nc.Dataset(sys.argv[2],"r")
        ncv=nc_file.variables
        x=ncv["x"][:]
        y=ncv["y"][:]
        x_inc=x[1]-x[0]
        y_inc=y[1]-y[0]
        x_range=(x[0]-x_inc/2.,x[-1]+x_inc/2.)
        y_range=(y[0]-y_inc/2.,y[-1]+y_inc/2.)
        img = rasterize(sys.argv[1], x_inc, "white", x_range, y_range )
    else:
        img = rasterize(sys.argv[1], 300, "white")

    img.save(sys.argv[1][:-4]+".bmp")