False Color Images

[bpartridge]

As I mentioned in the Inference group conference call, for debugging/paper purposes I've duplicated the SDSS's false-color image generation pipeline in Python. It can create images based on either a catalog or the loaded pixel data, if passed a Tractor as an argument. Here are examples from loaded pixel data and an inferred catalog from our MCMC work, respectively:

I need to clean up the code and comments, and fix it to work with Tractor instances that have more than 3 Images, but I'll post the code here in the meantime. If someone else wants to improve it or take a stab at it, feel free!

def falseColorImage(tractor, useModel=False):
  """
  Returns a PIL image which can be saved with rval.save(fname)
  or shown with matplotlib.pyplot.imshow(rval).

  Mimics the SDSS3 photoop pipeline stage sdss_frame_jpg.pro,
  which in turn calls djs_rgb_make.pro and polywarp_shift.pro.

  Outputs a false color image by mapping the i, r, and g bands
  to the RGB color channels, respectively.
  Because the images can have slight offsets due to the camera
  properties, the i and g bands are shifted so that they
  have the same centers as the r band.
  In the original pipeline, RGB channels are saturated at 30.
  (contiguous saturated pixels are replaced with their average 
  color), then scaled by 0.4 * [5., 6.5, 9.], then put through
  a nonlinear function [r,g,b] *= asinh(N*(r+g+b))/(N*(r+g+b))
  where N = 2.25.

  We simplify this by removing the nonlinearity and saturation.

  (Interesting tidbit from code: it seems there are 1361 unique
  rows in each image.)

  For now, we assume there are only images for the specified bands.
  """

  refimg = getFirstImgWithBand(tractor, 'r')

  # sdss_frame_jpg.pro
  # which does a polywarp shift of degree 1 (unless "keyword" set?)
  # filling with constant 0
  from scipy.ndimage.interpolation import shift

  def imgToCorrectedData(img):
    if useModel:
      data = tractor.getModelImage(img)
    else:
      data = img.getImage()

    return shiftBandImageToRef(img, refimg, data=data)

  idata = (0.4 * 5.) * imgToCorrectedData(getFirstImgWithBand(tractor, 'i'))
  rdata = (0.4 * 6.5) * imgToCorrectedData(refimg)
  gdata = (0.4 * 9.) * imgToCorrectedData(getFirstImgWithBand(tractor, 'g'))

  # print np.max(gdata)

  stacked = np.dstack((idata, rdata, gdata))

  stacked = np.flipud(stacked)

  # Perform saturation on individual pixels
  stacked[stacked < 0.] = 0.
  stacked[stacked > 30.] = 30.

  nonlin = 2.25
  sum_stacked = np.sum(stacked, axis=2)
  nonlin_mult = np.arcsinh(nonlin * sum_stacked) / (nonlin * sum_stacked + 1e-7)
  stacked *= nonlin_mult[:,:,np.newaxis]

  # Perform RGB saturation, which should be minimal
  stacked[stacked > 1.] = 1.

  from scipy.misc import toimage
  return toimage(stacked / np.max(stacked) * 256.0, channel_axis=2)

def getFirstImgWithBand(tractor, bandname):
  for img in tractor.getImages():
    if getBandNameForImage(img) == bandname:
      return img
  raise Exception("Tractor must have an image in band %s" % bandname)

def shiftBandImageToRef(img, refimg, data=None):
  if data is None:
    data = img.getImage()

  refpixctr = [n/2 for n in refimg.getShape()]
  refpos = refimg.getWcs().pixelToPosition(*refpixctr)

  from scipy.ndimage.interpolation import shift

  assert img.getShape() == refimg.getShape()

  pixdiff = img.getWcs().positionToPixel(refpos)
  pixshift = (refpixctr[1]-pixdiff[1], refpixctr[0]-pixdiff[0])

  shifted = shift(data, pixshift, order=1,
                  mode='constant', cval=0.0, prefilter=False)

  return shifted