graph: fix PNG download button

tensorboard/plugins/graph/tf_graph/tf-graph-scene.ts

@@ -27,6 +27,7 @@ import * as tf_graph_scene from '../tf_graph_common/scene';
 import * as tf_graph_scene_node from '../tf_graph_common/node';
 import * as tf_graph_util from '../tf_graph_common/util';
 import * as tf_graph_layout from '../tf_graph_common/layout';
+import * as tf_graph_minimap from '../tf_graph_common/minimap';
 import * as tf_graph_render from '../tf_graph_common/render';
 import {template} from './tf-graph-scene.html';
 
@@ -126,10 +127,8 @@ class TfGraphScene2
 
   /**
    * A minimap object to notify for zoom events.
-   * This property is a tf.scene.Minimap object.
    */
-  @property({type: Object})
-  minimap: object;
+  private minimap: tf_graph_minimap.Minimap;
 
   /*
    * Dictionary for easily stylizing nodes when state changes.
@@ -497,6 +496,9 @@ class TfGraphScene2
       }.bind(this)
     );
   }
+  getImageBlob(): Promise<Blob> {
+    return this.minimap.getImageBlob();
+  }
   isNodeSelected(n) {
     return n === this.selectedNode;
   }

tensorboard/plugins/graph/tf_graph/tf-graph.ts

@@ -332,6 +332,9 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
   fit() {
     (this.$.scene as any).fit();
   }
+  getImageBlob(): Promise<Blob> {
+    return (this.$.scene as any).getImageBlob();
+  }
   _graphChanged() {
     if (!this.graphHierarchy) {
       return;

tensorboard/plugins/graph/tf_graph_app/tf-graph-app.ts

@@ -98,6 +98,7 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
             on-fit-tap="_fit"
             trace-inputs="{{_traceInputs}}"
             auto-extract-nodes="{{_autoExtractNodes}}"
+            on-download-image-requested="_onDownloadImageRequested"
           ></tf-graph-controls>
           <tf-graph-loader
             id="loader"
@@ -201,4 +202,7 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
   _fit() {
     (this.$$('#graphboard') as any).fit();
   }
+  _onDownloadImageRequested(filename: string) {
+    (this.$$('#graphboard') as any).downloadAsImage(filename);
+  }
 }

tensorboard/plugins/graph/tf_graph_board/tf-graph-board.ts

@@ -295,6 +295,14 @@ class TfGraphBoard extends LegacyElementMixin(PolymerElement) {
   fit() {
     (this.$.graph as any).fit();
   }
+  async downloadAsImage(filename: string) {
+    const blob = await (this.$.graph as any).getImageBlob();
+    const element = document.createElement('a');
+    (element as any).href = (URL as any).createObjectURL(blob);
+    element.download = filename;
+    element.click();
+    URL.revokeObjectURL(element.href);
+  }
   /** True if the progress is not complete yet (< 100 %). */
   _isNotComplete(progress) {
     return progress.value < 100;

tensorboard/plugins/graph/tf_graph_common/minimap.ts

@@ -22,7 +22,6 @@ export class Minimap {
   private canvas: HTMLCanvasElement;
   /** A buffer canvas used for temporary drawing to avoid flickering. */
   private canvasBuffer: HTMLCanvasElement;
-  private download: HTMLLinkElement;
   private downloadCanvas: HTMLCanvasElement;
   /** The minimap svg used for holding the viewpoint rectangle. */
   private minimapSvg: SVGSVGElement;
@@ -141,6 +140,16 @@ export class Minimap {
       d3.zoomIdentity.translate(mainX, mainY).scale(this.scaleMain)
     );
   }
+  /**
+   * Takes a snapshot of the graph's image as a Blob.
+   */
+  getImageBlob(): Promise<Blob> {
+    return new Promise<Blob>((resolve) => {
+      this.downloadCanvas.toBlob((blob) => {
+        resolve(blob);
+      }, 'image/png');
+    });
+  }
   /**
    * Redraws the minimap. Should be called whenever the main svg
    * was updated (e.g. when a node was expanded).
@@ -159,31 +168,6 @@ export class Minimap {
       // detached from the dom.
       return;
     }
-    let $download = d3.select('#graphdownload');
-    this.download = <HTMLLinkElement>$download.node();
-    $download.on('click', (d) => {
-      // Revoke the old URL, if any. Then, generate a new URL.
-      URL.revokeObjectURL(this.download.href);
-      // We can't use the `HTMLCanvasElement.toBlob` API because it does
-      // not have a synchronous variant, and we need to update this href
-      // synchronously. Instead, we create a blob manually from the data
-      // URL.
-      const dataUrl = this.downloadCanvas.toDataURL('image/png');
-      const prefix = dataUrl.slice(0, dataUrl.indexOf(','));
-      if (!prefix.endsWith(';base64')) {
-        console.warn(
-          `non-base64 data URL (${prefix}); cannot use blob download`
-        );
-        (this.download as any).href = dataUrl;
-        return;
-      }
-      const data = atob(dataUrl.slice(dataUrl.indexOf(',') + 1));
-      const bytes = new Uint8Array(data.length).map((_, i) =>
-        data.charCodeAt(i)
-      );
-      const blob = new Blob([bytes], {type: 'image/png'});
-      (this.download as any).href = (URL as any).createObjectURL(blob);
-    });
     let $svg = d3.select(this.svg);
     // Read all the style rules in the document and embed them into the svg.
     // The svg needs to be self contained, i.e. all the style rules need to be

tensorboard/plugins/graph/tf_graph_controls/tf-graph-controls.ts

@@ -376,7 +376,6 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
           <iron-icon icon="file-download" class="button-icon"></iron-icon>
           <span class="button-text">Download PNG</span>
         </paper-button>
-        <a href="#" id="graphdownload" class="title" download="graph.png"></a>
       </div>
       <div class="control-holder runs">
         <div class="title">
@@ -1131,6 +1130,8 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
   })
   _legendOpened: boolean = true;
 
+  _downloadFilename = 'graph.png';
+
   _onGraphTypeChangedByUserGesture() {
     tf_graph_util.notifyDebugEvent({
       actionId: tb_debug.GraphDebugEventId.GRAPH_TYPE_CHANGED,
@@ -1315,7 +1316,7 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
     };
   }
   download() {
-    (this.$.graphdownload as HTMLElement).click();
+    this.fire('download-image-requested', this._downloadFilename);
   }
   _updateFileInput(e: Event) {
     const file = (e.target as HTMLInputElement).files[0];
@@ -1343,6 +1344,7 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
       // Select the first dataset by default.
       this._selectedRunIndex = 0;
     }
+    this._setDownloadFilename(this.datasets[this._selectedRunIndex]?.name);
   }
   _computeSelection(
     datasets: Dataset,
@@ -1369,9 +1371,7 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
     this._selectedTagIndex = 0;
     this._selectedGraphType = this._getDefaultSelectionType();
     this.traceInputs = false; // Set trace input to off-state.
-    this._setDownloadFilename(
-      this.datasets[runIndex] ? this.datasets[runIndex].name : ''
-    );
+    this._setDownloadFilename(this.datasets[runIndex]?.name);
   }
   _selectedTagIndexChanged(): void {
     this._selectedGraphType = this._getDefaultSelectionType();
@@ -1398,11 +1398,8 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
   _getFile() {
     (this.$$('#file') as HTMLElement).click();
   }
-  _setDownloadFilename(name: string) {
-    (this.$.graphdownload as HTMLElement).setAttribute(
-      'download',
-      name + '.png'
-    );
+  _setDownloadFilename(name?: string) {
+    this._downloadFilename = (name || 'graph') + '.png';
   }
   _statsNotNull(stats: tf_graph_proto.StepStats) {
     return stats !== null;

tensorboard/plugins/graph/tf_graph_dashboard/tf-graph-dashboard.ts

@@ -88,6 +88,7 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
         on-fit-tap="_fit"
         trace-inputs="{{_traceInputs}}"
         auto-extract-nodes="{{_autoExtractNodes}}"
+        on-download-image-requested="_onDownloadImageRequested"
       ></tf-graph-controls>
       <div
         class$="center [[_getGraphDisplayClassName(_selectedFile, _datasets)]]"
@@ -347,6 +348,9 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
   _fit() {
     (this.$$('#graphboard') as any).fit();
   }
+  _onDownloadImageRequested(event: CustomEvent) {
+    (this.$$('#graphboard') as any).downloadAsImage(event.detail as string);
+  }
   _getGraphDisplayClassName(_selectedFile: any, _datasets: any[]) {
     const isDataValid = _selectedFile || _datasets.length;
     return isDataValid ? '' : 'no-graph';