#!/usr/bin/env python

# writeRGB.py: two ways of saving out an 8-bit RGB PNG
# Gordon Kindlmann
#
# version 0.2
#  fixed bug in line 43: should be "out_data[j,i,2]"
# version 0.1, Mon Oct 11 13:30 CDT 2010
#  initial version

import Image, numpy

# First example, in which output data is created as an array
# of a particular size
out_data =  numpy.zeros((32,256,3),dtype=numpy.uint8)

# Example of filling in values with linear ramp
# In the assignment, colors in the output colormapped image will come from your function
for i in range(256):
  out_data[:,i,0] = i
  out_data[:,i,1] = 256-i
  out_data[:,i,2] = i

out_image = Image.fromarray(out_data,mode='RGB')
out_image.save("rgb1.png")

############################################

# Second example, creating the output numpy array based on shape
# of given image
# http://people.cs.uchicago.edu/~glk/237/proj2help/six.png
in_image = Image.open("six.png")
in_data = numpy.asarray(in_image)
(sy, sx) = (in_data.shape[0], in_data.shape[1])
out_data = numpy.zeros((sy, sx, 3),dtype=numpy.uint8)

# Again, your program must determine output colors based on 
# the colormap you create
for j in range(sy):
  for i in range(sx):
    out_data[j,i,0] = in_data[j,i]/256
    out_data[j,i,1] = 255 - in_data[j,i]/256
    out_data[j,i,2] = in_data[j,i]/256

out_image = Image.fromarray(out_data, "RGB")
out_image.save("rgb2.png")