Support RGBA in imshow

imshow_rgba.patch

commit f37df59507f77cd9e7cad33c27b6c63dd7aae0e5
Author: Parsa Ghadimi <[email protected]>
Date:   Tue Mar 6 22:15:36 2018 +0330

    [Fix #196] Support RGB in imshow

diff --git a/src/im.ts b/src/im.ts
index adbd34e..26f0c55 100644
--- a/src/im.ts
+++ b/src/im.ts
@@ -16,10 +16,17 @@
 // This module allows Propel to read PNG and JPG images.
 
 import { fill, Tensor, tensor } from "./api";
+import { imshow } from "./matplotlib";
 import { convert } from "./tensor";
 import { Mode, TypedArray } from "./types";
 import { IS_NODE, nodeRequire } from "./util";
 
+export type Uint8Image = {
+  height: number,
+  width: number,
+  data: Uint8Array
+};
+
 function toTensor(data: Uint8Array, height: number, width: number,
                   mode: Mode): Tensor {
   let image = convert(data).reshape([height, width, 4]);
@@ -36,7 +43,7 @@ function toTensor(data: Uint8Array, height: number, width: number,
   throw new Error("Unsupported convertion mode.");
 }
 
-export function toUint8Image(image: Tensor) {
+export function toUint8Image(image: Tensor): Uint8Image {
   const shape = image.shape;
   let data: TypedArray;
   let height: number;
@@ -231,5 +238,5 @@ export async function imsave(tensor: Tensor,
     }
     throw new Error(`Unsupported image format "${handler}"`);
   }
-  // todo: show image instead of saving it
+  imshow(tensor);
 }
diff --git a/src/matplotlib.ts b/src/matplotlib.ts
index 456d76f..5002b90 100644
--- a/src/matplotlib.ts
+++ b/src/matplotlib.ts
@@ -14,6 +14,7 @@
  */
 
 import { Tensor } from "./api";
+import { toUint8Image } from "./im";
 import { ImshowData, OutputHandler, PlotData } from "./output_handler";
 import { assertEqual } from "./tensor_util";
 
@@ -72,14 +73,5 @@ export function imshow(tensor: Tensor): void {
     return;
   }
 
-  // Assuming image shape is [3, height, width] for RGB.
-  // [height, width] for monochrome.
-  assertEqual(tensor.shape.length, 2, "Assuming monochrome for now");
-  const data: ImshowData = {
-    height: tensor.shape[0],
-    width: tensor.shape[1],
-    values: Array.from(tensor.dataSync())
-  };
-
-  currentOutputHandler.imshow(data);
+  currentOutputHandler.imshow(toUint8Image(tensor));
 }
diff --git a/src/output_handler.ts b/src/output_handler.ts
index 5230e7d..bdc113a 100644
--- a/src/output_handler.ts
+++ b/src/output_handler.ts
@@ -15,13 +15,13 @@
 // handler can be loaded without pulling in the entire math library.
 
 import * as d3 from "d3";
+import { Uint8Image } from "./im";
 
 export type PlotData = Array<Array<{ x: number, y: number }>>;
-export type ImshowData = { height: number, width: number, values: number[] };
 
 export interface OutputHandler {
   plot(data: PlotData): void;
-  imshow(data: ImshowData): void;
+  imshow(data: Uint8Image): void;
 }
 
 export class OutputHandlerDOM implements OutputHandler {
@@ -136,23 +136,13 @@ export class OutputHandlerDOM implements OutputHandler {
       });
   }
 
-  imshow({ width, height, values }: ImshowData): void {
+  imshow({ width, height, data }: Uint8Image): void {
     const canvas = document.createElement("canvas");
     canvas.height = height;
     canvas.width = width;
     const ctx = canvas.getContext("2d");
-    const imageData = ctx.getImageData(0, 0, width, height);
-    const data = imageData.data;
-    for (let y = 0; y < height; ++y) {
-      for (let x = 0; x < width; ++x) {
-        let index = (y * width + x) * 4;
-        const value = values[y * width + x];
-        data[index]   = value; // red
-        data[++index] = value; // green
-        data[++index] = value; // blue
-        data[++index] = 255;   // alpha
-      }
-    }
+    const clampedArray = Uint8ClampedArray.from(data);
+    const imageData = new ImageData(clampedArray, width, height);
     ctx.putImageData(imageData, 0, 0);
     this.element.appendChild(canvas);
   }