Add visualization code

borglab · Nov 26, 2020 · 6ed20e4 · 6ed20e4
1 parent f3e371f
commit 6ed20e4
Show file tree

Hide file tree

Showing 4 changed files with 194 additions and 3 deletions.
diff --git a/Scripts/Fan04.swift b/Scripts/Fan04.swift
@@ -50,7 +50,7 @@ struct Fan04: ParsableCommand {
       }
     }()
 
-    let (fig, _, _) = runProbabilisticTracker(
+    let (fig, track, gt) = runProbabilisticTracker(
       directory: dataDir,
       encoder: pca,
       onTrack: trackId, forFrames: trackLength, withSampling: true,
@@ -60,5 +60,15 @@ struct Fan04: ParsableCommand {
 
     /// Actual track v.s. ground truth track
     fig.savefig("Results/fan04/fan04_track\(trackId)_\(featureSize).pdf", bbox_inches: "tight")
+    fig.savefig("Results/fan04/fan04_track\(trackId)_\(featureSize).png", bbox_inches: "tight")
+
+    let json = JSONEncoder()
+    json.outputFormatting = .prettyPrinted
+
+    let track_data = try! json.encode(track)
+    try! track_data.write(to: URL(fileURLWithPath: "Results/fan04/fan04_track_\(trackId)_\(featureSize).json"))
+
+    let gt_data = try! json.encode(gt)
+    try! gt_data.write(to: URL(fileURLWithPath: "Results/fan04/fan04_gt_\(trackId)_\(featureSize).json"))
   }
 }
diff --git a/Scripts/Fan05.swift b/Scripts/Fan05.swift
@@ -0,0 +1,79 @@
+import ArgumentParser
+
+import SwiftFusion
+import BeeDataset
+import BeeTracking
+import PythonKit
+import Foundation
+import TensorFlow
+
+/// Fan05: Error Landscape
+struct Fan05: ParsableCommand {
+  @Option(help: "Size of feature space")
+  var featureSize: Int = 100
+
+  @Option(help: "Which frame to show")
+  var frameId: Int = 0
+
+  @Option(help: "Which track to show")
+  var trackId: Int = 0
+
+  // Visualize error landscape of PCA
+  // Make sure you have a folder `Results/fan05` before running
+  func run() {
+    let dataDir = URL(fileURLWithPath: "./OIST_Data")
+
+    let np = Python.import("numpy")
+
+      // train foreground and background model and create tracker
+    let trainingData = OISTBeeVideo(directory: dataDir, length: 100)!
+    // let testData = OISTBeeVideo(directory: dataDir, afterIndex: 100, length: forFrames)!
+    let (imageHeight, imageWidth, imageChannels) = (40, 70, 1)
+    // let encoder = RandomProjection(fromShape: TensorShape([imageHeight, imageWidth, imageChannels]), toFeatureSize: featureSize)
+
+    // let encoder = PCAEncoder(
+    //   withBasis: Tensor<Double>(numpy: np.load("./pca_U_\(featureSize).npy"))!,
+    //   andMean: Tensor<Double>(numpy: np.load("./pca_mu_\(featureSize).npy"))!
+    // )
+    var encoder = DenseRAE(
+      imageHeight: imageHeight, imageWidth: imageWidth, imageChannels: imageChannels,
+      hiddenDimension: 100, latentDimension: featureSize
+    )
+
+    encoder.load(weights: np.load("./oist_rae_weight_\(featureSize).npy", allow_pickle: true))
+
+    let (fg, bg, statistics) = getTrainingBatches(
+      dataset: trainingData, boundingBoxSize: (40, 70),
+      fgBatchSize: 3000,
+      bgBatchSize: 3000,
+      fgRandomFrameCount: 100,
+      bgRandomFrameCount: 100,
+      useCache: true
+    )
+
+    let batchPositive = encoder.encode(fg)
+    let foregroundModel = MultivariateGaussian(from: batchPositive, regularizer: 1e-3)
+
+    let batchNegative = encoder.encode(bg)
+    let backgroundModel = MultivariateGaussian(from: batchNegative, regularizer: 1e-3)
+
+    let deltaXRange = Array(-60..<60).map { Double($0) }
+    let deltaYRange = Array(-40..<40).map { Double($0) }
+
+    let datasetToShow = OISTBeeVideo(directory: dataDir, afterIndex: frameId - 1, length: 2)!
+    let frame = datasetToShow.frames[1]
+    let pose = datasetToShow.tracks[trackId].boxes[0].center
+    let (fig, _) = plotErrorPlaneTranslation(
+      frame: frame,
+      at: pose,
+      deltaXs: deltaXRange,
+      deltaYs: deltaYRange,
+      statistics: statistics,
+      encoder: encoder,
+      foregroundModel: foregroundModel,
+      backgroundModel: backgroundModel
+    )
+    fig.savefig("Results/fan05/fan05_pf_ae_mg_mg_\(trackId)_\(frameId)_\(featureSize).pdf", bbox_inches: "tight")
+    fig.savefig("Results/fan05/fan05_pf_ae_mg_mg_\(trackId)_\(frameId)_\(featureSize).png", bbox_inches: "tight")
+  }
+}
diff --git a/Scripts/main.swift b/Scripts/main.swift
@@ -17,7 +17,7 @@ import PenguinParallelWithFoundation
 
 struct Scripts: ParsableCommand {
   static var configuration = CommandConfiguration(
-    subcommands: [Fan01.self, Fan02.self, Fan03.self, Fan04.self,
+    subcommands: [Fan01.self, Fan02.self, Fan03.self, Fan04.self, Fan05.self,
                   Frank01.self, Frank02.self, Frank03.self, Frank04.self])
 }
 

diff --git a/Sources/BeeTracking/Visualizations.swift b/Sources/BeeTracking/Visualizations.swift
@@ -69,4 +69,106 @@ public func plotPatchWithGT(frame: Tensor<Float>, actual: Pose2, expected: Pose2
   ).makeNumpyArray() / 255.0)
   ax[1].title.set_text("Labeling")
   return (fig, ax)
-}
+}
+
+public func errorPlaneTranslation<
+  Encoder: AppearanceModelEncoder,
+  FGModel: GenerativeDensity,
+  BGModel: GenerativeDensity
+> (
+  frame: Tensor<Float>,
+  at: Pose2,
+  deltaXs: [Double],
+  deltaYs: [Double],
+  statistics: FrameStatistics,
+  encoder: Encoder,
+  foregroundModel: FGModel,
+  backgroundModel: BGModel
+) -> (fg: Tensor<Double>, bg: Tensor<Double>, e: Tensor<Double>) {
+  let targetSize = (40, 70)
+  let measurement = statistics.normalized(frame)
+
+  func error(_ pose: Pose2) -> (fg: Double, bg: Double, e: Double) {
+    let region = OrientedBoundingBox(center: pose, rows: targetSize.0, cols: targetSize.1)
+    let patch = Tensor<Double>(measurement.patch(at: region, outputSize: targetSize))
+    let features = encoder.encode(patch.expandingShape(at: 0)).squeezingShape(at: 0)
+
+    let fg_nll = foregroundModel.negativeLogLikelihood(features)
+    let bg_nll = backgroundModel.negativeLogLikelihood(features)
+    let result = fg_nll - bg_nll
+
+    /// TODO: What is the idiomatic way of avoiding negative probability here?
+    return (fg: fg_nll, bg: bg_nll, e: result)
+  }
+
+  var fg = Tensor<Double>(zeros: [deltaYs.count, deltaXs.count])
+  var bg = Tensor<Double>(zeros: [deltaYs.count, deltaXs.count])
+  var errors = Tensor<Double>(zeros: [deltaYs.count, deltaXs.count])
+  for (i, dx) in deltaXs.enumerated() {
+    for (j, dy) in deltaYs.enumerated() {
+      let (fg_nll, bg_nll, e) = error(at * Pose2(dx, dy, 0.0))
+      /// x is horiz movement, but is vertical dim in imshow
+      fg[j, i] = Tensor(fg_nll)
+      bg[j, i] = Tensor(bg_nll)
+      errors[j, i] = Tensor(e)
+    }
+  }
+  return (fg, bg, errors)
+}
+
+public func plotErrorPlaneTranslation<
+  Encoder: AppearanceModelEncoder,
+  FGModel: GenerativeDensity,
+  BGModel: GenerativeDensity
+> (
+  frame: Tensor<Float>,
+  at pose: Pose2,
+  deltaXs: [Double],
+  deltaYs: [Double],
+  statistics: FrameStatistics,
+  encoder: Encoder,
+  foregroundModel: FGModel,
+  backgroundModel: BGModel
+) -> (PythonObject, PythonObject) {
+  let plt = Python.import("matplotlib.pyplot")
+  let (fg, bg, e) = errorPlaneTranslation(
+    frame: frame,
+    at: pose,
+    deltaXs: deltaXs,
+    deltaYs: deltaYs,
+    statistics: statistics,
+    encoder: encoder,
+    foregroundModel: foregroundModel,
+    backgroundModel: backgroundModel
+  )
+
+  // let trans_mins = [e, fg, bg].map { $0.min() }
+  // let trans_maxs = [e, fg, bg].map { $0.max() }
+  // let trans_min = Tensor<Double>(trans_mins).min().scalarized()
+  // let trans_max = Tensor<Double>(trans_maxs).max().scalarized()
+  // let targetSize = (40, 70)
+  let (fig, axs) = plt.subplots(2, 2, figsize: Python.tuple([12, 10])).tuple2
+  let img_m = axs[0][0].imshow(frame.patch(
+    at: OrientedBoundingBox(center: pose, rows: 40 + 20 * 2, cols: 70 + 20 * 2)
+  ).makeNumpyArray() / 255.0, cmap: "gray")
+  fig.colorbar(img_m, ax: axs[0][0])
+  axs[0][0].title.set_text("Image")
+  axs[0][0].set(xlabel: "x displacement", ylabel: "y displacement")
+
+  let fg_m = axs[0][1].imshow(fg.makeNumpyArray(), cmap: "hot", interpolation: "nearest") // , vmin: trans_min, vmax: trans_max)
+  fig.colorbar(fg_m, ax: axs[0][1])
+  axs[0][1].title.set_text("Foreground Response")
+  axs[0][1].set(xlabel: "x displacement", ylabel: "y displacement")
+
+  let bg_m = axs[1][0].imshow(bg.makeNumpyArray(), cmap: "hot", interpolation: "nearest" ) // , vmin: trans_min, vmax: trans_max)
+  fig.colorbar(bg_m, ax: axs[1][0])
+  axs[1][0].title.set_text("Background Response")
+  axs[1][0].set(xlabel: "x displacement", ylabel: "y displacement")
+
+  let pcm = axs[1][1].imshow(e.makeNumpyArray(), cmap: "hot", interpolation: "nearest") // , vmin: trans_min, vmax: trans_max)
+  fig.colorbar(pcm, ax: axs[1][1])
+  axs[1][1].title.set_text("Total Response")
+  axs[1][1].set(xlabel: "x displacement", ylabel: "y displacement")
+
+  return (fig, axs)
+}