--- a/testing/marionette/driver.js
+++ b/testing/marionette/driver.js
@@ -2403,19 +2403,26 @@ GeckoDriver.prototype.clearImportedScrip
  *     string.
  */
 GeckoDriver.prototype.takeScreenshot = function (cmd, resp) {
   let {id, highlights, full, hash} = cmd.parameters;
   highlights = highlights || [];
 
   switch (this.context) {
     case Context.CHROME:
-      let canvas = capture.viewport(this.getCurrentWindow());
-      resp.body.value = capture.toBase64(canvas);
-      break;
+      let container = {frame: this.getCurrentWindow()};
+      let highlightEls = [];
+
+      for (let h of highlights) {
+        let el = this.curBrowser.seenEls.get(h, container);
+        highlightEls.push(el);
+      }
+
+      let canvas = capture.viewport(this.getCurrentWindow(), highlightEls);
+      return capture.toBase64(canvas);
 
     case Context.CONTENT:
       if (hash) {
         return this.listener.getScreenshotHash(id, full, highlights);
       } else {
         return this.listener.takeScreenshot(id, full, highlights);
       }
   }