Factorie decode using MPLP

factorie_MPLP_decode.scala

/**
* Construct Factor graph and run MAP
*/
def decode(unaryPot: DenseMatrix[Double],
           pairwisePot: DenseMatrix[Double],
           adj: AdjacencyList): Array[Label] = {

  val nSuperpixels = unaryPot.cols
  val nClasses = unaryPot.rows

  assert(nClasses == NUM_CLASSES)
  assert(pairwisePot.rows == NUM_CLASSES)

  object PixelDomain extends DiscreteDomain(nClasses)

  class Pixel(i: Int) extends DiscreteVariable(i) {
    def domain = PixelDomain
  }

  def getUnaryFactor(yi: Pixel, idx: Int): Factor = {
    new Factor1(yi) {
      def score(k: Pixel#Value) = -unaryPot(k.intValue, idx)
    }
  }

  def getPairwiseFactor(yi: Pixel, yj: Pixel): Factor = {
    new Factor2(yi, yj) {
      def score(i: Pixel#Value, j: Pixel#Value) = pairwisePot(i.intValue, j.intValue)
    }
  }

  val pixelSeq: IndexedSeq[Pixel] =
    (0 until nSuperpixels).map(x => new Pixel(12))

  val unaryFactors: IndexedSeq[Factor] =
    (0 until nSuperpixels).map {
      case idx =>
        getUnaryFactor(pixelSeq(idx), idx)
    }

  val pairwiseFactors =
    (0 until nSuperpixels).flatMap {
      case thisIdx =>
        val thisFactors = new ArrayBuffer[Factor]

        adj(thisIdx).foreach {
          case nextIdx =>
            thisFactors ++=
              getPairwiseFactor(pixelSeq(thisIdx), pixelSeq(nextIdx))
        }
        thisFactors
    }

  val model = new ItemizedModel
  model ++= unaryFactors
  model ++= pairwiseFactors

  val maxIterations = 300
  val maximizer = new MaximizeByMPLP(maxIterations)
  val assgn = maximizer.infer(pixelSeq, model).mapAssignment

  val mapLabels: Array[Label] = (0 until nSuperpixels).map {
    idx =>
      assgn(pixelSeq(idx)).intValue
  }.toArray

  mapLabels
}