/*
 * Interactive sketch to demonstrate chromakey masking.
 * Author: Michael Goldwasser
 *
 * Images taken from:
 * http://upload.wikimedia.org/wikipedia/commons/c/c2/Girl_in_front_of_a_green_background.jpg
 * http://dvcreators.net/images/front_shot.jpg
 *
 * We allow the user to separately control:
 * -- The chroma key color. This is done by clicking on a pixel of either reference image
 * -- The function used for computing relative difference between two colors:
 *      pressing 'M' (the default) gets you Manhatten distance in RGB space
 *      pressing 'E' gets you Euclidean distance in RGB space
 *      pressing 'L' gets you Euclidean distance in L*a*b space 
 * -- The threshold value used for masking can be changed as follows
 *    For the left image, UP increases and DOWN decreases threshold
 *    For the right image, RIGHT increases and LEFT decreases threshold
 *
 * Finally, if you like what you see, 'S' saves the current screen to file.
 */

// only needed for running javascript version in browswer
/* @pjs preload="exampleGreen.jpg,exampleBlue.jpg"; */

PImage[] original = new PImage[2];
color[] chromakey = new color[2];
float[] threshold = new float[2];
float threshIncr = 2.5;
char diffChoice = 'M';   // should be E(uclidean), M(anhatten), or L(*a*b)

// support for showing placeholder image while computing
boolean showPlaceholder = false;
int[] updateStatus = { 0, 0 };

int w;  // width of each of the panels
int h;  // height of canvas

float s = 0.65;  // use this to downsize window for smaller display

void setup() {
  //original[0] = loadImage("exampleBlue.jpg");
  //original[1] = loadImage("exampleGreen.jpg");
  original[0] = loadImage("http://mathcs.slu.edu/~goldwasser/144/activity/masking/exampleBlue.jpg");
  original[1] = loadImage("http://mathcs.slu.edu/~goldwasser/144/activity/masking/exampleGreen.jpg");

  w = max(original[0].width, original[1].width);
  h = max(original[0].height, original[1].height);
  if (s == 0.5) {
    size(1284, 480);  // hardcode for javascript/browswer version only
  } else {
    size(int(s*4*w), int(s*h));
  }
  background(0);
  stroke(0);
  fill(0);
  textAlign(LEFT, BOTTOM);

  // draw originals
  pushMatrix();
  scale(s, s);
  image(original[0], 0, 0);
  image(original[1], 2*w, 0);
  popMatrix();

  // initialize parameters
  for (int j=0; j <= 1; j++) {
    chromakey[j] = original[j].get(0, 0);  // use top-left pixel as default
    threshold[j] = 45;                     // somewhat arbitrary choice
    updateStatus[j] = 1;                   // force initial refresh
  }
} // end setup()


// redisplay the given image in the given panel (or a temporary message if img is null)
void displayPanel(int j, PImage img) {
  pushMatrix();
  scale(s, s);
  int left = w*(1+2*j);
  fill(255*(1-j));            // white for image 0, black for image 1
  rect(left, 0, w, h);
  fill(255, 0, 0);
  if (img != null) {
    image(img, left, 0);
    text(display(j), left+5, h);
  } else {
    text("Recomputing...", left+5, h);
  }
  popMatrix();
}

// recompute and display masked version of image j given current settings
void recreateMask(int j) {
  // make copy so as to leave original unchanged
  PImage temp = original[j].get(0, 0, original[j].width, original[j].height);
  chromaMask(temp, chromakey[j], threshold[j]);
  displayPanel(j, temp);
}


// Modify the given image by masking all pixels that are
// an approximate match for the given key color.
void chromaMask(PImage img, color keyColor, float tolerance) {
  PImage msk = createImage(img.width, img.height, RGB); //the mask should be an image the same size as the original
  for (int y=0; y<img.height; y++) {
    for (int x=0; x<img.width; x++) {
      color c = img.get(x, y);
      float diff = computeDiff(c, keyColor);
      if (diff <= tolerance) {
        msk.set(x, y, color(0));
      } else {
        msk.set(x, y, color(255));
      }
    }
  }
  img.mask(msk);
}

// Generic function to manage choice of color difference functions
float computeDiff(color c1, color c2) {
  if (diffChoice == 'M') {
    return manhattenDiff(c1, c2);
  } else if (diffChoice == 'L') {
    return labDiff(c1, c2);
  } else {  // default
    return euclideanDiff(c1, c2);
  }
}

// This version is based on computing Euclidean distance in RGB space
// (although this is not a particular good choice of measure)
float euclideanDiff(color c1, color c2) {
  return sqrt(sq(red(c1) - red(c2)) + sq(green(c1) - green(c2)) + sq(blue(c1) - blue(c2)));
}

float manhattenDiff(color c1, color c2) {
  return abs(red(c1)-red(c2)) + abs(green(c1)-green(c2)) + abs(blue(c1)-blue(c2));
}

// This version is based on computing Euclidean distance in L*a*b space
// see http://www.easyrgb.com/index.php?X=MATH for formula
float[] rgbToLab(color c) {
  // first go to XYZ space
  float[] v = new float[3];
  v[0] = red(c)/255.0;
  v[1] = green(c)/255.0;
  v[2] = blue(c)/255.0;
  for (int j=0; j < 3; j++) {
    if (v[j] > 0.0405) {
      v[j] = pow((v[j]+0.055)/1.055, 2.4);
    } else {
      v[j] /= 12.92;
    }
    v[j] *= 100;
  }
  float x = v[0] * 0.4124 + v[1] * 0.3576 + v[2] * 0.1805;
  float y = v[0] * 0.2126 + v[1] * 0.7152 + v[2] * 0.0722;
  float z = v[0] * 0.0193 + v[1] * 0.1192 + v[2] * 0.9505;

  // now convert from XYX to L*a*b
  v[0] = x / 95.047;
  v[1] = y / 100.000;
  v[2] = z / 108.883;

  for (int j=0; j < 3; j++) {
    if (v[j] > 0.00856) {
      v[j] = pow(v[j], 1.0/3.0);
    } else {
      v[j] = 7.787 * v[j] + 16.0/116;
    }
  }

  float[] lab = { 
    (116 * v[1]) - 16, 500 * (v[0]-v[1]), 200 * (v[1] - v[2])
    };
    return lab;
}

float labDiff(color c1, color c2) {
  float[] lab1 = rgbToLab(c1);
  float[] lab2 = rgbToLab(c2);
  return sqrt(sq(lab1[0]-lab2[0])+sq(lab1[1]-lab2[1])+sq(lab1[2]-lab2[2]));
}

String display(int j) {
  String diff = "";
  switch (diffChoice) {
    case ('E'):
    diff = "RGB (Euclidean)";
    break;
    case ('M'):
    diff = "RGB (Manhatten)";
    break;
    case ('L'):
    diff = "CIE-Lab";
    break;
  }
  String output = "Diff: " + diff + "\nkey ";
  color c = chromakey[j];
  output += "("+red(c)+","+green(c)+","+blue(c)+")";
  output += "\nthreshold="+nf(threshold[j], 0, 1);
  return output;
}

// User interactions to adjust key colors and thresholds
void draw() {
  for (int j=0; j < 2; j++) {
    if (updateStatus[j] == 1) {
      // add "Recomputing..." label while processing
      if (showPlaceholder) {
        displayPanel(j, null);
      }
      updateStatus[j]++;
    } else if (updateStatus[j] == 2) {
      updateStatus[j] = 0;
      recreateMask(j);
    }
  }
}

void mouseClicked() {
  int x = int(mouseX / s);
  int panel = x / w;
  if (panel % 2 == 0) {  // even-numbered panels are originals
    int j = panel / 2;
    int adjustedX = x - panel * w;
    chromakey[j] = original[j].get(adjustedX, int(mouseY/s));
    updateStatus[j] = 1;
  }
}

void keyPressed() {
  if (keyCode == UP) {
    threshold[0] += threshIncr;
    updateStatus[0] = 1;
  } else if (keyCode == DOWN) {
    threshold[0] -= threshIncr;
    if (threshold[0] < 0) {
      threshold[0] = 0;
    }
    updateStatus[0] = 1;
  } else if (keyCode == RIGHT) {
    threshold[1] += threshIncr;
    updateStatus[1] = 1;
  } else if (keyCode == LEFT) {
    threshold[1] -= threshIncr;
    if (threshold[1] < 0) {
      threshold[1] = 0;
    }
    updateStatus[1] = 1;
  }

  if (key == 'E' || key == 'M' || key == 'L') {
    if (key != diffChoice) {
      diffChoice = key;
      updateStatus[0] = 1;
      updateStatus[1] = 1;
    }
  }

  if (key == 'S') {
    save("chromaKey.jpg");
  }
}