'''BuildImageCatalog.py

1. Start in the indicated folder.
2. Initialize a file whose name is the folder name, followed by "-catalog.txt"
3. Find all the folder's subfolders.
4. In each folder, get a list of the JPG images.
5. For each image in the list:
    a. load it into PixelMath;
    b. compute a set of features.
    c. write an entry in the file with this format:
        IMG_NO: <int>; FOLDER: <folder name>; SUBFOLDER: <subfolder name>; FILENAME: <filename>; MEAN_COLOR: <RGB value>; MODE_COLOR: <RGB value>; BLUE_FRAC: <float>; BEIGE_FRAC: <float>; GREEN_FRAC: <float>; HORIZ_EDGE_FRAC: <float>; VERT_EDGE_FRAC: <float>; OMNIDIR_EDGE_FRAC: <float>; WIDTH: <int>; HEIGHT: <int>
    d. print progress info, such as a subset of the items in (c).
6. Close the file, and report on the time taken.
'''

import os
from math import floor


def list_JPGs(folder):
  files = os.listdir(folder)
  return filter(lambda f: f[-4:]==".JPG", files)

def list_folders(folder):
  files = os.listdir(folder)
  return filter(lambda f: os.path.isdir(os.path.join(folder,f)), files)

IMG_NO = 0

def process_image(source_path, info):
  print "Preparing to process the image: "+source_path
  pmOpenImage(1, source_path)
  width = pmGetImageWidth(1)
  height = pmGetImageHeight(1)
  pmSetSource1(1)
  newWin = pmNewImage(2, "Processed version of "+source_path, width, height, 255, 0, 0)
  pmSetDestination(newWin)
  # HERE IS WHERE TO PUT THE CODE TO REALLY PROCESS IMAGES:

  # MEAN_COLOR: <RGB value>; 
  npixels = width*height
  avgRed   = pmTotalRed(1)/npixels
  avgGreen = pmTotalGreen(1)/npixels
  avgBlue  = pmTotalBlue(1)/npixels
  info += "; MEAN_COLOR: "+str((avgRed, avgGreen, avgBlue))

  # MODE_COLOR: <RGB value>; 
  # First quantize colors more coarsely (3 bits of red, 3 of green, 2 of blue)
  # The commented-out formula that follows produces a version of the image that
  # gives an idea of the distortion implied by the quantization
  #pmSetFormula("(floor(S1(x,y)/RGB(32,32,64))*RGB(32,32,64)")
  # The actual quantization puts all three of RGB and by into one bye, and
  # since that's done in all three of R,G, and B, we happen to get a monochrome image.
  # We'll only really use the Red component for histogramming (any of the 3 will do).
  pmSetFormula("floor(Red1(x,y)/32)*32 + floor(Green1(x,y)/32)*8 + floor(Blue1(x,y)/64)")
  pmCompute()
  # Next create a histogram of these values.
  h = pmHistogramRed(newWin)
  print "Histogram is: "+str(h)
  # Find mode:
  maxSoFar = -1
  modeVal = 0
  for i in range(len(h)):
    elt = h[i]
    if elt > maxSoFar: 
      maxSoFar = elt
      modeVal = i
  mode_color = (floor(modeVal/32)*32, floor((modeVal % 32)/4)*32, (modeVal % 4)*64)
  print "The mode seems to be ... "+str(mode_color)
  info += "; MODE_COLOR: "+str(mode_color)

  # Hue-based fractions. First we need to compute Hue vals and vals that indicate acceptable sat. and val.
  # BLUE_FRAC: <float>; 
  pmSetFormula("if Hue1(x,y)>0.5 and Hue1(x,y)<0.7 and Sat1(x,y)>0.1 then rgb(1,0,255) else 0")
  pmCompute()
  totalBlue = pmTotalRed(newWin) # Count up the pixels having the right hue and sat profile.
  blue_frac = float(totalBlue) / npixels
  info += "; BLUE_FRAC: "+str(blue_frac)

  # BEIGE_FRAC: <float>; 
  pmSetFormula("if Hue1(x,y)>0.09 and Hue1(x,y)<0.13 and Sat1(x,y)>0.1 and Val1(x,y)>0.5 then rgb(1,0,255) else 0")
  pmCompute()
  totalBeige = pmTotalRed(newWin) # Count up the pixels having the right hue and sat profile.
  beige_frac = float(totalBeige) / npixels
  info += "; BEIGE_FRAC: "+str(beige_frac)

  # GREEN_FRAC: <float>; 
  pmSetFormula("if Hue1(x,y)>0.14 and Hue1(x,y)<0.4 and Sat1(x,y)>0.1 and Val1(x,y)>0.1 then rgb(1,0,255) else 0")
  pmCompute()
  totalGreen = pmTotalRed(newWin) # Count up the pixels having the right hue and sat profile.
  green_frac = float(totalGreen) / npixels
  info += "; GREEN_FRAC: "+str(green_frac)

  # Compute edge characteristics, making use of the Sobel gradients gx and gy.
  # Let's put gx in the destination image's R component, gy in the G component, and
  # the angle they imply in B.
  R_formula = "127+(S1(x-1,y-1)+2*S1(x-1,y)+S1(x-1,y+1)-S1(x+1,y-1)-2*S1(x+1,y)-S1(x+1,y+1))/4"
  G_formula = "127+(S1(x-1,y-1)+2*S1(x,y-1)+S1(x+1,y-1)-S1(x-1,y+1)-2*S1(x,y+1)-S1(x+1,y+1))/4"
  B_formula = "Angle(Red2(x,y)-127,Green2(x,y)-127) * 30"
  pmSetFormula("RGB("+R_formula+","+G_formula+","+B_formula+")")
  pmEnableColors(1, 1, 0) # Don't bother with blue yet.
  pmCompute()
  # New compute angle from GX and GY (taking from R and G and putting into B)
  pmEnableColors(0, 0, 1) # New compute only blue (angle)
  pmSetSource2(newWin)    # Source and Dest. the same for this step.
  pmCompute()
  pmEnableColors(1, 1, 1) # Restore normal enablement.
  # Now detect vertical and horizontal edge pixels.
  edgepix = pmNewImage(3, "Edge pixels", width, height, 0, 0, 100)
  pmSetSource1(newWin)
  pmSetDestination(edgepix)
  #pmSetFormula("if (sqr(red1(x,y)-127)+sqr(green1(x,y)-127))>400 and Blue1(x,y)/30 < rgb(pi, pi/2, 2*pi)+0.1 and Blue1(x,y)/30 > rgb(pi, pi/2, 0)-0.1 then rgb(255, 255, 255) else 0")
  pmSetFormula("if (sqr(red1(x,y)-127)+sqr(green1(x,y)-127))>400 then "+\
    "rgb(if {Blue1(x,y)/30 < pi+0.1 and Blue1(x,y)/30 > pi-0.1} or {Blue1(x,y)/30 < 0.1} or {Blue1(x,y)/30 > 2*pi-0.1} then 1 else 0, "+\
        "if {Blue1(x,y)/30 < pi/2 +0.1 and Blue1(x,y)/30 > pi/2 -0.1} or {Blue1(x,y)/30 < 3*pi/2 +0.1 and Blue1(x,y)/30 > 3*pi/2 -0.1} then 1 else 0, 1) else 0")
  pmCompute()
  totalEdges = pmTotalBlue(edgepix) 
  # HORIZ_EDGE_FRAC: <float>; 
  totalHoriz = pmTotalRed(edgepix) 
  if totalEdges==0:
    horiz_frac = 0.0
  else:
    horiz_frac = float(totalHoriz)/totalEdges
  info += ": HORIZ_EDGE_FRAC: "+str(horiz_frac)

  # VERT_EDGE_FRAC: <float>
  totalVert = pmTotalGreen(edgepix) 
  if totalEdges==0:
    vert_frac = 0.0
  else:
    vert_frac = float(totalVert)/totalEdges
  info += ": VERT_EDGE_FRAC: "+str(vert_frac)

  # OMNIDIR_EDGE_FRAC: <float>
  edge_frac = float(totalEdges)/npixels
  info += ": OMNIDIR_EDGE_FRAC: "+str(edge_frac)

  info += "; WIDTH: "+str(width)
  info += "; HEIGHT: "+str(height)
  return info

def process_images_in_folder(parent_folder, current_folder, catalog_file_handle):
  # go into the source folder and process all the JPGs there.
  global IMG_NO
  new_source_path = os.path.join(parent_folder, current_folder)
  print "new_source_path = "+new_source_path
  JPGs = list_JPGs(new_source_path)
  for fn in JPGs:
    print "Processing "+fn
    IMG_NO += 1
    info = "IMG_NO: "+str(IMG_NO)+"; FOLDER: "+parent_folder
    info += "; SUBFOLDER: "+current_folder+"; FILENAME: "+fn
    new_info = process_image(os.path.join(new_source_path, fn), info)+"\n"
    print "Writing into the catalog: "+new_info
    catalog_file_handle.write(new_info)

def build_catalog(folder_name):
  # create catalog file.
  try:
    cat_file = open(folder_name+"-catalog.txt","w")
  except:
    print "Could not open the catalog file for writing."
    return

  global IMG_NO
  IMG_NO = 0
  print "Requesting list of folders in "+folder_name
  folders = list_folders(folder_name)
  print "Folders to be processed: "+str(folders)
  for f in folders:
    process_images_in_folder(folder_name, f, cat_file)
  cat_file.close()
  print "Done creating catalog."

import time
start_time = time.time()
build_catalog('Images-for-Pattern-Rec_reduced_by_0.1') # Do the work.
end_time = time.time()
print "Elapsed time: "+str(end_time - start_time)