SketchRNN (#189)

* init sketchrnn

* strokes

* add models

* stylefixes

* some more bugfixes

* comments

* update dependencies

Author: reiinakano

package-lock.json

@@ -86,9 +86,10 @@
     ]
   },
   "dependencies": {
-    "@tensorflow-models/mobilenet": "0.1.0",
-    "@tensorflow-models/posenet": "0.1.3",
-    "@tensorflow/tfjs": "0.11.4",
+    "@magenta/sketch": "^0.1.2",
+    "@tensorflow-models/mobilenet": "0.2.2",
+    "@tensorflow-models/posenet": "0.2.2",
+    "@tensorflow/tfjs": "0.13.0",
     "events": "^3.0.0"
   }
 }

package.json

@@ -0,0 +1,68 @@
+// Copyright (c) 2018 ml5
+//
+// This software is released under the MIT License.
+// https://opensource.org/licenses/MIT
+
+/* eslint prefer-destructuring: ['error', {AssignmentExpression: {array: false}}] */
+/* eslint no-await-in-loop: 'off' */
+/*
+SketchRNN
+*/
+
+import * as ms from '@magenta/sketch';
+import callCallback from '../utils/callcallback';
+import modelPaths from './models';
+
+const PATH_START_LARGE = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/large_models/';
+const PATH_START_SMALL = 'https://storage.googleapis.com/quickdraw-models/sketchRNN/models/';
+const PATH_END = '.gen.json';
+
+class SketchRNN {
+  constructor(model, callback, large = true) {
+    let checkpointUrl = model;
+    if (modelPaths.has(checkpointUrl)) {
+      checkpointUrl = (large ? PATH_START_LARGE : PATH_START_SMALL) + checkpointUrl + PATH_END;
+    }
+    this.defaults = {
+      temperature: 0.65,
+      pixelFactor: 3.0,
+    };
+    this.model = new ms.SketchRNN(checkpointUrl);
+    this.penState = this.model.zeroInput();
+    this.ready = callCallback(this.model.initialize(), callback);
+  }
+
+  async generateInternal(options, strokes) {
+    const temperature = +options.temperature || this.defaults.temperature;
+    const pixelFactor = +options.pixelFactor || this.defaults.pixelFactor;
+
+    await this.ready;
+    if (!this.rnnState) {
+      this.rnnState = this.model.zeroState();
+      this.model.setPixelFactor(pixelFactor);
+    }
+
+    if (Array.isArray(strokes) && strokes.length) {
+      this.rnnState = this.model.updateStrokes(strokes, this.rnnState);
+    }
+    this.rnnState = this.model.update(this.penState, this.rnnState);
+    const pdf = this.model.getPDF(this.rnnState, temperature);
+    this.penState = this.model.sample(pdf);
+    return this.penState;
+  }
+
+  async generate(options, strokes, callback) {
+    return callCallback(this.generateInternal(options, strokes), callback);
+  }
+
+  reset() {
+    this.penState = this.model.zeroInput();
+    if (this.rnnState) {
+      this.rnnState = this.model.zeroState();
+    }
+  }
+}
+
+const SketchRNNGenerator = (model, callback, large = true) => new SketchRNN(model, callback, large);
+
+export default SketchRNNGenerator;

src/SketchRNN/index.js

@@ -0,0 +1,120 @@
+const models = [
+  'alarm_clock',
+  'ambulance',
+  'angel',
+  'ant',
+  'antyoga',
+  'backpack',
+  'barn',
+  'basket',
+  'bear',
+  'bee',
+  'beeflower',
+  'bicycle',
+  'bird',
+  'book',
+  'brain',
+  'bridge',
+  'bulldozer',
+  'bus',
+  'butterfly',
+  'cactus',
+  'calendar',
+  'castle',
+  'cat',
+  'catbus',
+  'catpig',
+  'chair',
+  'couch',
+  'crab',
+  'crabchair',
+  'crabrabbitfacepig',
+  'cruise_ship',
+  'diving_board',
+  'dog',
+  'dogbunny',
+  'dolphin',
+  'duck',
+  'elephant',
+  'elephantpig',
+  'eye',
+  'face',
+  'fan',
+  'fire_hydrant',
+  'firetruck',
+  'flamingo',
+  'flower',
+  'floweryoga',
+  'frog',
+  'frogsofa',
+  'garden',
+  'hand',
+  'hedgeberry',
+  'hedgehog',
+  'helicopter',
+  'kangaroo',
+  'key',
+  'lantern',
+  'lighthouse',
+  'lion',
+  'lionsheep',
+  'lobster',
+  'map',
+  'mermaid',
+  'monapassport',
+  'monkey',
+  'mosquito',
+  'octopus',
+  'owl',
+  'paintbrush',
+  'palm_tree',
+  'parrot',
+  'passport',
+  'peas',
+  'penguin',
+  'pig',
+  'pigsheep',
+  'pineapple',
+  'pool',
+  'postcard',
+  'power_outlet',
+  'rabbit',
+  'rabbitturtle',
+  'radio',
+  'radioface',
+  'rain',
+  'rhinoceros',
+  'rifle',
+  'roller_coaster',
+  'sandwich',
+  'scorpion',
+  'sea_turtle',
+  'sheep',
+  'skull',
+  'snail',
+  'snowflake',
+  'speedboat',
+  'spider',
+  'squirrel',
+  'steak',
+  'stove',
+  'strawberry',
+  'swan',
+  'swing_set',
+  'the_mona_lisa',
+  'tiger',
+  'toothbrush',
+  'toothpaste',
+  'tractor',
+  'trombone',
+  'truck',
+  'whale',
+  'windmill',
+  'yoga',
+  'yogabicycle',
+  'everything',
+];
+
+const modelPaths = new Set(models);
+
+export default modelPaths;

src/SketchRNN/models.js

@@ -14,6 +14,7 @@ import * as imageUtils from './utils/imageUtilities';
 import styleTransfer from './StyleTransfer/';
 import LSTMGenerator from './LSTM/';
 import pix2pix from './Pix2pix/';
+import SketchRNN from './SketchRNN';
 
 module.exports = {
   imageClassifier,
@@ -25,6 +26,7 @@ module.exports = {
   poseNet,
   LSTMGenerator,
   pix2pix,
+  SketchRNN,
   ...imageUtils,
   tf,
 };