/*
 * Interactive sketch to demonstrate chromakey masking.
 * Author: Michael Goldwasser
 *
 * 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 with UP/DOWN arrows
 * -- The image source can be chosen from a collection by typing the
 *    digit that corresponds to its index in the collection (with 0 = default)
 *
 * Finally, if you like what you see, 'S' saves the current screen to file.
 */

int desiredHeight = 400;
int footerHeight = 30;   // for extra info

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

// enter up to 10 strings to use as input
String collection[] = {
  "http://cs.slu.edu/~goldwasser/1050/lectures/chromakey/exampleBlue.jpg",  // http://dvcreators.net/wp-content/uploads/sites/11/images/front_shot.jpg
  "http://cs.slu.edu/~goldwasser/1050/lectures/chromakey/exampleGreen.jpg", // originally http://upload.wikimedia.org/wikipedia/commons/c/c2/Girl_in_front_of_a_green_background.jpg
  "http://ohsnapsphotobooth.com/wp-content/uploads/2010/09/Green-screen-puppets.jpg",
  "http://cdn2.hubspot.net/hub/310020/file-442100700-jpg/images/green-color-in-regalia-thumbnail.jpg?t=1449152719391&width=150",
  "http://image.gala.de/v1/cms/MT/ewan-mcgregor-17dez-ge_5393214-original-portrait_6col.jpg?v=4084629",
  "http://ecx.images-amazon.com/images/I/81RyJJGWGmL._SL1500_.jpg",
  "http://www.glamourcraft.com/school/images/spring_ck0.jpg",
  "https://i.ytimg.com/vi/_-dTC8EH5LQ/maxresdefault.jpg",
  "http://www.supermanhomepage.com/images/superman-returns8/ew-greenscreen.jpg",
};

PImage original;
PImage mask;
PImage masked;
PGraphics composed;
boolean needRefresh;
color chromakey;
float threshold = 45;    // somewhat arbitrary choice
float threshIncr = 2.5;
char diffChoice = 'M';   // should be E(uclidean), M(anhatten), or L(*a*b)

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


void setup() {
  noStroke();
  textAlign(LEFT, BOTTOM);
  surface.setResizable(true);
} // end setup()

void selectImage(int j) {
  original = loadImage(collection[j]);
  w = original.width;
  h = original.height;
  s = 1.0*(desiredHeight-footerHeight)/h;
  surface.setSize(int(4*s*w), desiredHeight); // four panels and a footer
  needRefresh = true;
  chromakey = original.get(0, 0);  // use top-left pixel as default key color
  mask = createImage(w,h,RGB);
  composed = createGraphics(w,h);
}

// User interactions to adjust key colors and thresholds
void draw() {
  if (frameCount == 1) {
    selectImage(0);
  }

  // do footer and parameter displays
  background(255);
  fill(chromakey);
  rect(0,desiredHeight-footerHeight/2,width,footerHeight/2);
  textSize(0.8*footerHeight/2);
  fill(0);
  text(configString(),20,desiredHeight-footerHeight/2);

  // scaled versions of the various views
  scale(s);
  fill(255, 0, 0);      // red for text
  image(original, 0, 0);
  if (needRefresh) {
    recompute();
  }
  image(mask, w, 0);
  image(masked, 2*w, 0);
  image(composed, 3*w, 0);
}

void mouseClicked() {
  int x = int(mouseX / s);  // adjust for scaling
  int panel = x / w;
  if (panel == 0) {
    chromakey = original.get(x, int(mouseY/s));
    needRefresh = true;
  }
}

void keyPressed() {
  if (keyCode == UP) {
    threshold += threshIncr;
    needRefresh = true;
  } else if (keyCode == DOWN) {
    threshold -= threshIncr;
    if (threshold < 0) {
      threshold = 0;
    }
    needRefresh = true;
  }

  // upcase character
  if (key >= 'a' && key <= 'z') {
    int adj = 'A' - 'a';
    key = char(key + adj);
  }

  if (key == 'E' || key == 'M' || key == 'L') {
    if (key != diffChoice) {
      diffChoice = key;
      needRefresh = true;;
    }
  }

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

  if (key >= '0' && key <= '9') {
    int j = key - '0';
    if (j < collection.length) {
      selectImage(j);
    }
  }
}


void recompute() {
  needRefresh = false;
  createMask();
  masked = original.get(0,0,w,h);
  masked.mask(mask);   // wow; that's an unfortunate choice of names
  
  // draw checkerboard
  float sz = 30;  // 30x30 grid squares
  composed.beginDraw();
  composed.fill(255,255,0);
  composed.rect(0,0,w,h);
  composed.fill(255,127,0);
  composed.stroke(255);
  for (int a=0; a*sz < w; a ++) {
    for (int b=0; b*sz < h; b ++) {
      if (a % 2 == b % 2) {
        composed.rect(a*sz, b*sz, sz, sz);
      }      
    }
  }
  composed.image(masked,0,0);
  composed.endDraw();
}


void createMask() {
  for (int x=0; x<w; x++) {
    for (int y=0; y<h; y++) {
      color c = original.get(x, y);
      float diff = computeDiff(c, chromakey);
      if (diff <= threshold) {
        mask.set(x, y, color(0));
      } else {
        mask.set(x, y, color(255));
      }
    }
  }
}

// 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 configString() {
  String sep = ", ";
  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 + sep + "key ";
  color c = chromakey;
  output += "("+red(c)+","+green(c)+","+blue(c)+")";
  output += sep + "threshold="+nf(threshold, 0, 1);
  return output;
}