From 3fb3fb76f16d81662248ec84c5513a379688d497 Mon Sep 17 00:00:00 2001
From: Jonathan Lipps <jlipps@gmail.com>
Date: Thu, 19 Jul 2018 13:49:04 +0900
Subject: [PATCH] add some jimp helper functionality

---
 lib/image-util.js            | 37 +++++++++++++++++++++++++++++++++++-
 lib/mjpeg.js                 | 14 +++-----------
 test/image-util-e2e-specs.js | 34 ++++++++++++++++++++++++++++++++-
 3 files changed, 72 insertions(+), 13 deletions(-)

diff --git a/lib/image-util.js b/lib/image-util.js
index 5053c8b2..5c7ad55e 100644
--- a/lib/image-util.js
+++ b/lib/image-util.js
@@ -1,8 +1,11 @@
 import _ from 'lodash';
+import Jimp from 'jimp';
 import { Buffer } from 'buffer';
 import { PNG } from 'pngjs';
 import B from 'bluebird';
 import { hasValue } from './util';
+
+const { MIME_JPEG, MIME_PNG, MIME_BMP } = Jimp;
 let cv = null;
 
 /**
@@ -52,6 +55,37 @@ const AVAILABLE_MATCHING_FUNCTIONS = [
   'BruteForceSL2',
 ];
 
+/**
+ * Utility function to get a Jimp image object from buffer or base64 data. Jimp
+ * is a great library however it does IO in the constructor so it's not
+ * convenient for our async/await model.
+ *
+ * @param {Buffer|string} data - binary image buffer or base64-encoded image
+ * string
+ */
+async function getJimpImage (data) {
+  return await new B((resolve, reject) => {
+    if (!_.isString(data) && !_.isBuffer(data)) {
+      throw new Error("Must initialize jimp object with string or buffer");
+    }
+    // if data is a string, assume it is a base64-encoded image
+    if (_.isString(data)) {
+      data = Buffer.from(data, 'base64');
+    }
+    new Jimp(data, (err, imgObj) => {
+      if (err) {
+        return reject(err);
+      }
+      if (!imgObj) {
+        return reject(new Error("Could not create jimp image from that data"));
+      }
+      imgObj._getBuffer = imgObj.getBuffer.bind(imgObj);
+      imgObj.getBuffer = B.promisify(imgObj._getBuffer, {context: imgObj});
+      resolve(imgObj);
+    });
+  });
+}
+
 /**
  * @throws {Error} If opencv4nodejs module is not installed or cannot be loaded
  */
@@ -510,4 +544,5 @@ function getRectIntersection (rect, imageSize) {
 }
 
 export { cropBase64Image, base64ToImage, imageToBase64, cropImage,
-         getImagesMatches, getImagesSimilarity, getImageOccurrence };
+         getImagesMatches, getImagesSimilarity, getImageOccurrence,
+         getJimpImage, MIME_JPEG, MIME_PNG, MIME_BMP };
diff --git a/lib/mjpeg.js b/lib/mjpeg.js
index 9a4bc551..a1e37035 100644
--- a/lib/mjpeg.js
+++ b/lib/mjpeg.js
@@ -3,7 +3,7 @@ import request from 'request';
 import log from './logger';
 import http from 'http';
 import B from 'bluebird';
-import Jimp from 'jimp';
+import { getJimpImage, MIME_PNG } from './image-util';
 import MJpegConsumer from 'mjpeg-consumer';
 import mJpegServer from 'mjpeg-server';
 import { Writable } from 'stream';
@@ -49,16 +49,8 @@ class MJpegStream extends Writable {
       return null;
     }
 
-    const jpg = await new B((res, rej) => {
-      new Jimp(this.lastChunk, (err, img) => {
-        if (err) {
-          return rej(err);
-        }
-        res(img);
-      });
-    });
-
-    return await B.promisify(jpg.getBuffer, {context: jpg})(Jimp.MIME_PNG);
+    const jpg = await getJimpImage(this.lastChunk);
+    return await jpg.getBuffer(MIME_PNG);
   }
 
   /**
diff --git a/test/image-util-e2e-specs.js b/test/image-util-e2e-specs.js
index a6a8580f..bfeed98d 100644
--- a/test/image-util-e2e-specs.js
+++ b/test/image-util-e2e-specs.js
@@ -1,6 +1,8 @@
 import { base64ToImage, imageToBase64, cropImage,
-         getImagesMatches, getImagesSimilarity, getImageOccurrence } from '../lib/image-util';
+         getImagesMatches, getImagesSimilarity, getImageOccurrence,
+         getJimpImage, MIME_PNG } from '../lib/image-util';
 import path from 'path';
+import _ from 'lodash';
 import chai from 'chai';
 import { fs } from 'appium-support';
 import chaiAsPromised from 'chai-as-promised';
@@ -129,4 +131,34 @@ describe('image-util', function () {
       });
     });
   });
+
+  describe('Jimp helpers', function () {
+    it('should get a jimp object using image buffer', async function () {
+      const base64Image = await getImage('cropped-image.b64');
+      const imageBuffer = Buffer.from(base64Image, 'base64');
+      const jimpImg = await getJimpImage(imageBuffer);
+      jimpImg.hash().should.eql('80000000000');
+      jimpImg.bitmap.height.should.eql(485);
+      jimpImg.bitmap.width.should.eql(323);
+    });
+    it('should get a jimp object using b64 string', async function () {
+      const base64Image = await getImage('cropped-image.b64');
+      const jimpImg = await getJimpImage(base64Image);
+      jimpImg.hash().should.eql('80000000000');
+      jimpImg.bitmap.height.should.eql(485);
+      jimpImg.bitmap.width.should.eql(323);
+    });
+    it('should error with incorrect data type', async function () {
+      await getJimpImage(1234).should.eventually.be.rejectedWith(/string or buffer/);
+    });
+    it('should error with incorrect image data', async function () {
+      await getJimpImage('foo').should.eventually.be.rejectedWith(/Could not create jimp image/);
+    });
+    it('should get an image buffer via the overridden getBuffer method', async function () {
+      const base64Image = await getImage('cropped-image.b64');
+      const jimpImg = await getJimpImage(base64Image);
+      const buf = await jimpImg.getBuffer(MIME_PNG);
+      _.isBuffer(buf).should.be.true;
+    });
+  });
 });