diff --git a/tfjs-automl/README.md b/tfjs-automl/README.md
index 91f0ff5f236..8f79d8f90f9 100644
--- a/tfjs-automl/README.md
+++ b/tfjs-automl/README.md
@@ -22,7 +22,7 @@ If you are using CDN:
We support the following types of AutoML Edge models:
1) [Image classification](#image-classification)
-2) **[In progress]** [Object detection](#object-detection)
+2) [Object detection](#object-detection)
## Image classification
@@ -67,7 +67,7 @@ a 3D [`Tensor`](https://js.tensorflow.org/api/latest/#class:Tensor):
```js
const img = document.getElementById('img');
-const options = {};
+const options = {centerCrop: true};
const predictions = await model.classify(img, options);
```
@@ -88,6 +88,90 @@ probabilities:
]
```
+### Advanced usage
+
+Advanced users can access the underlying
+[`GraphModel`](https://js.tensorflow.org/api/latest/#class:GraphModel) via
+`model.graphModel`. The `GraphModel` allows users to call lower level methods
+such as `predict()`, `execute()` and `executeAsync()` which return tensors.
+
+`model.dictionary` gives you access to the ordered list of labels.
+
## Object detection
-TODO(smilkov): Write this when object detection is ready.
+AutoML Object detection model will output the following set of files:
+- `model.json`, the model topology
+- `dict.txt`, a newline-separated list of labels
+- One or more of `*.bin` files which hold the weights
+
+Make sure you can access those files as static assets from your web app by serving them locally or on Google Cloud Storage.
+
+### Demo
+
+The object detection demo lives in
+[demo/object_classification](./demo/object_classification). To run it:
+
+```sh
+cd demo/object_detection
+yarn
+yarn watch
+```
+
+This will start a local HTTP server on port 1234 that serves the demo.
+
+### Loading the model
+```js
+import * as automl from '@tensorflow/tfjs-automl';
+const modelUrl = 'model.json'; // URL to the model.json file.
+const model = await automl.loadObjectDetection(modelUrl);
+```
+
+### Making a prediction
+The input `img` can be
+[`HTMLImageElement`](https://developer.mozilla.org/en-US/docs/Web/API/HTMLImageElement),
+[`HTMLCanvasElement`](https://developer.mozilla.org/en-US/docs/Web/API/HTMLCanvasElement),
+[`HTMLVideoElement`](https://developer.mozilla.org/en-US/docs/Web/API/HTMLVideoElement),
+[`ImageData`](https://developer.mozilla.org/en-US/docs/Web/API/ImageData) or
+a 3D [`Tensor`](https://js.tensorflow.org/api/latest/#class:Tensor):
+
+```html
+AutoML Object detection demo
+
+
+
+
+
+
+
+
+
{
+ Promise {
options = sanitizeOptions(options);
const scores = tidy(() => {
@@ -60,10 +70,6 @@ export async function loadImageClassification(modelUrl: string):
return new ImageClassificationModel(model, dict);
}
-function imageToTensor(img: ImageInput): Tensor3D {
- return img instanceof Tensor ? img : browser.fromPixels(img);
-}
-
function sanitizeOptions(options: ImageClassificationOptions) {
options = options || {} as ImageClassificationOptions;
if (options.centerCrop == null) {
@@ -72,15 +78,6 @@ function sanitizeOptions(options: ImageClassificationOptions) {
return options;
}
-/** Loads and parses the dictionary. */
-async function loadDictionary(modelUrl: string): Promise {
- const prefixUrl = modelUrl.slice(0, modelUrl.lastIndexOf('/'));
- const dictUrl = `${prefixUrl}/dict.txt`;
- const response = await util.fetch(dictUrl);
- const text = await response.text();
- return text.trim().split('\n');
-}
-
/** Center crops an image */
function centerCropAndResize(img: Tensor3D) {
return tidy(() => {
diff --git a/tfjs-automl/src/img_classification_test.ts b/tfjs-automl/src/img_classification_test.ts
index 4c7096a407b..be215288697 100644
--- a/tfjs-automl/src/img_classification_test.ts
+++ b/tfjs-automl/src/img_classification_test.ts
@@ -15,12 +15,12 @@
* =============================================================================
*/
+import {GraphModel} from '@tensorflow/tfjs-converter';
import * as tf from '@tensorflow/tfjs-core';
-import {Tensor3D, test_util} from '@tensorflow/tfjs-core';
import {BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
import * as automl from './index';
-import {ClassificationPrediction} from './types';
+import {fetchImage} from './test_util';
const MODEL_URL =
'https://storage.googleapis.com/tfjs-testing/tfjs-automl/img_classification/model.json';
@@ -28,7 +28,7 @@ const MODEL_URL =
const DAISY_URL =
'https://storage.googleapis.com/tfjs-testing/tfjs-automl/img_classification/daisy.jpg';
-describeWithFlags('nodejs+browser integration', {}, () => {
+describeWithFlags('image classification', {}, () => {
let model: automl.ImageClassificationModel = null;
beforeAll(async () => {
@@ -36,39 +36,54 @@ describeWithFlags('nodejs+browser integration', {}, () => {
});
it('make prediction from a tensor', async () => {
- const img: Tensor3D = tf.zeros([100, 80, 3]);
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
const predictions = await model.classify(img);
expect(predictions[0].label).toBe('daisy');
expect(predictions[1].label).toBe('dandelion');
expect(predictions[2].label).toBe('roses');
- test_util.expectNumbersClose(predictions[0].prob, 0.5806022);
- test_util.expectNumbersClose(predictions[1].prob, 0.32249659);
- test_util.expectNumbersClose(predictions[2].prob, 0.0283515);
+ tf.test_util.expectNumbersClose(predictions[0].prob, 0.5806022);
+ tf.test_util.expectNumbersClose(predictions[1].prob, 0.32249659);
+ tf.test_util.expectNumbersClose(predictions[2].prob, 0.0283515);
});
it('make prediction from a tensor without cropping', async () => {
- const img: Tensor3D = tf.zeros([100, 80, 3]);
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
const predictions = await model.classify(img, {centerCrop: false});
expect(predictions[0].label).toBe('daisy');
expect(predictions[1].label).toBe('dandelion');
expect(predictions[2].label).toBe('roses');
- test_util.expectNumbersClose(predictions[0].prob, 0.5806022);
- test_util.expectNumbersClose(predictions[1].prob, 0.32249659);
- test_util.expectNumbersClose(predictions[2].prob, 0.0283515);
+ tf.test_util.expectNumbersClose(predictions[0].prob, 0.5806022);
+ tf.test_util.expectNumbersClose(predictions[1].prob, 0.32249659);
+ tf.test_util.expectNumbersClose(predictions[2].prob, 0.0283515);
});
it('no memory leak when making a prediction', async () => {
- const img: Tensor3D = tf.zeros([100, 80, 3]);
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
const numTensorsBefore = tf.memory().numTensors;
await model.classify(img);
const numTensorsAfter = tf.memory().numTensors;
expect(numTensorsAfter).toEqual(numTensorsBefore);
});
+
+ it('has access to dictionary', () => {
+ expect(model.dictionary).toEqual([
+ 'daisy', 'dandelion', 'roses', 'sunflowers', 'tulips'
+ ]);
+ });
+
+ it('can access the underlying graph model', () => {
+ expect(model.graphModel instanceof GraphModel).toBe(true);
+ expect(model.graphModel.inputNodes).toEqual(['image']);
+ expect(model.graphModel.outputNodes).toEqual(['scores']);
+ const img: tf.Tensor = tf.zeros([1, 224, 224, 3]);
+ const scores = model.graphModel.predict(img) as tf.Tensor;
+ expect(scores.shape).toEqual([1, 5]);
+ });
});
-describeWithFlags('browser integration', BROWSER_ENVS, () => {
+describeWithFlags('image classification browser', BROWSER_ENVS, () => {
let model: automl.ImageClassificationModel = null;
let daisyImg: HTMLImageElement;
@@ -78,7 +93,7 @@ describeWithFlags('browser integration', BROWSER_ENVS, () => {
});
function assertTop3PredsForDaisy(
- predictions: ClassificationPrediction, centerCrop: boolean) {
+ predictions: automl.ImagePrediction[], centerCrop: boolean) {
const probs = centerCrop ? [0.9310929, 0.0273733, 0.0130559] :
[0.8411523, 0.0729438, 0.03020708];
expect(predictions[0].label).toBe('daisy');
@@ -113,20 +128,3 @@ describeWithFlags('browser integration', BROWSER_ENVS, () => {
assertTop3PredsForDaisy(predictions, false /* centerCrop */);
});
});
-
-async function fetchImage(url: string): Promise {
- const response = await fetch(url);
- const blob = await response.blob();
- const img = new Image();
- const blobUrl = URL.createObjectURL(blob);
- return new Promise((resolve, reject) => {
- img.onload = () => {
- URL.revokeObjectURL(blobUrl);
- resolve(img);
- };
- img.onerror = (evt /* Arg is an event, not error. Can't rethrow it */) => {
- reject(new Error('Failed to load blob as image.'));
- };
- img.src = blobUrl;
- });
-}
diff --git a/tfjs-automl/src/index.ts b/tfjs-automl/src/index.ts
index 233b6630e53..a93c3b7d8dd 100644
--- a/tfjs-automl/src/index.ts
+++ b/tfjs-automl/src/index.ts
@@ -16,7 +16,10 @@
*/
// Image classification API.
-export {ImageClassificationModel, loadImageClassification} from './img_classification';
-export {ClassificationPrediction, ImageClassificationOptions, ImageInput} from './types';
+export {ImageClassificationModel, ImageClassificationOptions, ImagePrediction, loadImageClassification} from './img_classification';
+// Object detection API.
+export {Box, loadObjectDetection, ObjectDetectionModel, ObjectDetectionOptions, PredictedObject} from './object_detection';
+// Shared API.
+export {ImageInput} from './types';
export {version} from './version';
diff --git a/tfjs-automl/src/object_detection.ts b/tfjs-automl/src/object_detection.ts
new file mode 100644
index 00000000000..7a1e38eeeb5
--- /dev/null
+++ b/tfjs-automl/src/object_detection.ts
@@ -0,0 +1,175 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GraphModel, loadGraphModel} from '@tensorflow/tfjs-converter';
+import {dispose, image, Tensor, Tensor2D, tidy} from '@tensorflow/tfjs-core';
+
+import {ImageInput} from './types';
+import {imageToTensor, loadDictionary} from './util';
+
+const DEFAULT_TOPK = 20;
+const DEFAULT_IOU_THRESHOLD = 0.5;
+const DEFAULT_SCORE_THRESHOLD = 0.5;
+
+// TODO(smilkov): Change the name of the input nodes once we get a new model
+// from Cloud.
+const INPUT_NODE_NAME = 'ToFloat';
+const OUTPUT_NODE_NAMES =
+ ['Postprocessor/convert_scores', 'Postprocessor/Decode/transpose_1'];
+
+export interface ObjectDetectionOptions {
+ /**
+ * Only the `topk` most likely objects are returned. The actual number of
+ * objects might be less than this number.
+ */
+ topk?: number;
+ /**
+ * Intersection over union threshold. IoU is a metric between 0 and 1 used to
+ * measure the overlap of two boxes. The predicted boxes will not overlap more
+ * than the specified threshold.
+ */
+ iou?: number;
+ /** Boxes with score lower than this threshold will be ignored. */
+ score?: number;
+}
+
+/** Contains the coordinates of a bounding box. */
+export interface Box {
+ /** Number of pixels from the top of the image (top padding). */
+ top: number;
+ /** Number of pixels from the left of the image (left padding). */
+ left: number;
+ /** The width of the box. */
+ width: number;
+ /** The height of the box. */
+ height: number;
+}
+
+/** The predicted object, which holds the score, label and bounding box. */
+export interface PredictedObject {
+ box: Box;
+ score: number;
+ label: string;
+}
+
+export class ObjectDetectionModel {
+ constructor(public graphModel: GraphModel, public dictionary: string[]) {}
+
+ async detect(input: ImageInput, options?: ObjectDetectionOptions):
+ Promise {
+ options = sanitizeOptions(options);
+ const img = tidy(() => this.preprocess(input, options));
+ const [height, width] = [img.shape[1], img.shape[2]];
+ const feedDict: {[name: string]: Tensor} = {};
+ feedDict[INPUT_NODE_NAME] = img;
+ const [scoresTensor, boxesTensor] =
+ await this.graphModel.executeAsync(feedDict, OUTPUT_NODE_NAMES) as
+ Tensor[];
+
+ const [, numBoxes, numClasses] = scoresTensor.shape;
+ const [scores, boxes] =
+ await Promise.all([scoresTensor.data(), boxesTensor.data()]);
+ const {boxScores, boxLabels} =
+ calculateMostLikelyLabels(scores as Float32Array, numBoxes, numClasses);
+
+ // Sort the boxes by score, ignoring overlapping boxes.
+ const selectedBoxesTensor = await image.nonMaxSuppressionAsync(
+ boxesTensor as Tensor2D, boxScores, options.topk, options.iou,
+ options.score);
+ const selectedBoxes = await selectedBoxesTensor.data() as Int32Array;
+ dispose([img, scoresTensor, boxesTensor, selectedBoxesTensor]);
+
+ const result = buildDetectedObjects(
+ width, height, boxes as Float32Array, boxScores, boxLabels,
+ selectedBoxes, this.dictionary);
+ return result;
+ }
+
+ private preprocess(input: ImageInput, options: ObjectDetectionOptions) {
+ return imageToTensor(input).expandDims().toFloat();
+ }
+}
+
+export async function loadObjectDetection(modelUrl: string):
+ Promise {
+ const [model, dict] =
+ await Promise.all([loadGraphModel(modelUrl), loadDictionary(modelUrl)]);
+ return new ObjectDetectionModel(model, dict);
+}
+
+function sanitizeOptions(options: ObjectDetectionOptions) {
+ options = options || {} as ObjectDetectionOptions;
+ if (options.topk == null) {
+ options.topk = DEFAULT_TOPK;
+ }
+ if (options.iou == null) {
+ options.iou = DEFAULT_IOU_THRESHOLD;
+ }
+ if (options.score == null) {
+ options.score = DEFAULT_SCORE_THRESHOLD;
+ }
+ return options;
+}
+
+function calculateMostLikelyLabels(
+ scores: Float32Array, numBoxes: number,
+ numClasses: number): {boxScores: number[], boxLabels: number[]} {
+ // Holds a score for each box.
+ const boxScores: number[] = [];
+ // Holds the label id for each box.
+ const boxLabels: number[] = [];
+ for (let i = 0; i < numBoxes; i++) {
+ let maxScore = Number.MIN_VALUE;
+ let mostLikelyLabel = -1;
+ for (let j = 0; j < numClasses; j++) {
+ const flatIndex = i * numClasses + j;
+ const score = scores[flatIndex];
+ if (score > maxScore) {
+ maxScore = scores[flatIndex];
+ mostLikelyLabel = j;
+ }
+ }
+ boxScores[i] = maxScore;
+ boxLabels[i] = mostLikelyLabel;
+ }
+ return {boxScores, boxLabels};
+}
+
+function buildDetectedObjects(
+ width: number, height: number, boxes: Float32Array, boxScores: number[],
+ boxLabels: number[], selectedBoxes: Int32Array,
+ dictionary: string[]): PredictedObject[] {
+ const objects: PredictedObject[] = [];
+ // Each 2d rectangle is fully described with 4 coordinates.
+ const numBoxCoords = 4;
+ for (let i = 0; i < selectedBoxes.length; i++) {
+ const boxIndex = selectedBoxes[i];
+ const [top, left, bottom, right] = Array.from(boxes.slice(
+ boxIndex * numBoxCoords, boxIndex * numBoxCoords + numBoxCoords));
+ objects.push({
+ box: {
+ left: left * width,
+ top: top * height,
+ width: (right - left) * width,
+ height: (bottom - top) * height,
+ },
+ label: dictionary[boxLabels[boxIndex]],
+ score: boxScores[boxIndex],
+ });
+ }
+ return objects;
+}
diff --git a/tfjs-automl/src/object_detection_test.ts b/tfjs-automl/src/object_detection_test.ts
new file mode 100644
index 00000000000..131699d163e
--- /dev/null
+++ b/tfjs-automl/src/object_detection_test.ts
@@ -0,0 +1,140 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GraphModel} from '@tensorflow/tfjs-converter';
+import * as tf from '@tensorflow/tfjs-core';
+import {BROWSER_ENVS, describeWithFlags} from '@tensorflow/tfjs-core/dist/jasmine_util';
+
+import * as automl from './index';
+import {fetchImage} from './test_util';
+
+const MODEL_URL =
+ 'https://storage.googleapis.com/tfjs-testing/tfjs-automl/object_detection/model.json';
+
+const SALAD_IMAGE =
+ 'https://storage.googleapis.com/tfjs-testing/tfjs-automl/object_detection/test_image.jpg';
+
+describeWithFlags('object detection', {}, () => {
+ let model: automl.ObjectDetectionModel = null;
+
+ beforeAll(async () => {
+ model = await automl.loadObjectDetection(MODEL_URL);
+ });
+
+ it('prediction from a zero tensor', async () => {
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
+ const predictions = await model.detect(img);
+ expect(predictions.length).toBe(0);
+ });
+
+ it('prediction from a zero tensor with score threshold', async () => {
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
+ const options = {score: 0.11};
+ const predictions = await model.detect(img, options);
+ expect(predictions.length).toBe(1);
+ const {box: {left, top, width, height}, label, score} = predictions[0];
+
+ expect(label).toBe('Salad');
+ tf.test_util.expectNumbersClose(score, 0.116391);
+ tf.test_util.expectNumbersClose(left, 26.6741156);
+ tf.test_util.expectNumbersClose(top, 94.003981);
+ tf.test_util.expectNumbersClose(width, 22.6277661);
+ tf.test_util.expectNumbersClose(height, 6.30063414);
+ });
+
+ it('prediction with iou threshold', async () => {
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
+ const options: automl.ObjectDetectionOptions = {score: 0.10};
+ options.iou = 0.9;
+ let predictions = await model.detect(img, options);
+ expect(predictions.length).toBe(9);
+
+ options.iou = 0.1;
+ predictions = await model.detect(img, options);
+ expect(predictions.length).toBe(1);
+ });
+
+ it('no memory leak when making a prediction', async () => {
+ const img: tf.Tensor3D = tf.zeros([100, 80, 3]);
+ const numTensorsBefore = tf.memory().numTensors;
+ await model.detect(img);
+ const numTensorsAfter = tf.memory().numTensors;
+ expect(numTensorsAfter).toEqual(numTensorsBefore);
+ });
+
+ it('has access to dictionary', () => {
+ expect(model.dictionary).toEqual([
+ 'background', 'Tomato', 'Seafood', 'Salad', 'Baked Goods', 'Cheese'
+ ]);
+ });
+
+ it('can access the underlying graph model', () => {
+ expect(model.graphModel instanceof GraphModel).toBe(true);
+ });
+});
+
+describeWithFlags('object detection browser', BROWSER_ENVS, () => {
+ let model: automl.ObjectDetectionModel = null;
+ let saladImg: HTMLImageElement;
+
+ beforeAll(async () => {
+ model = await automl.loadObjectDetection(MODEL_URL);
+ saladImg = await fetchImage(SALAD_IMAGE);
+ });
+
+ function assertTop3PredsForSalad(predictions: automl.PredictedObject[]) {
+ expect(predictions.length).toBe(3);
+ const [top1, top2, top3] = predictions;
+ expect(top1.label).toBe('Tomato');
+ tf.test_util.expectNumbersClose(top1.score, 0.97170084);
+
+ expect(top2.label).toBe('Tomato');
+ tf.test_util.expectNumbersClose(top2.score, 0.93456619);
+
+ expect(top3.label).toBe('Salad');
+ tf.test_util.expectNumbersClose(top3.score, 0.9074271);
+ }
+
+ it('make prediction from an image element', async () => {
+ const predictions = await model.detect(saladImg);
+ assertTop3PredsForSalad(predictions);
+ });
+
+ it('make prediction from a canvas element', async () => {
+ // Copy the pixels from the image to a canvas.
+ const canvas = document.createElement('canvas');
+ const ctx = canvas.getContext('2d');
+ canvas.width = saladImg.width;
+ canvas.height = saladImg.height;
+ ctx.drawImage(saladImg, 0, 0, saladImg.width, saladImg.height);
+
+ const predictions = await model.detect(canvas);
+ assertTop3PredsForSalad(predictions);
+ });
+
+ it('make prediction from image data', async () => {
+ const canvas = document.createElement('canvas');
+ const ctx = canvas.getContext('2d');
+ canvas.width = saladImg.width;
+ canvas.height = saladImg.height;
+ ctx.drawImage(saladImg, 0, 0, saladImg.width, saladImg.height);
+ const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
+
+ const predictions = await model.detect(imageData);
+ assertTop3PredsForSalad(predictions);
+ });
+});
diff --git a/tfjs-automl/src/test_browser.ts b/tfjs-automl/src/test_browser.ts
index 27a978498b2..2b3c66bf163 100644
--- a/tfjs-automl/src/test_browser.ts
+++ b/tfjs-automl/src/test_browser.ts
@@ -17,6 +17,9 @@
import {setTestEnvs} from '@tensorflow/tfjs-core/dist/jasmine_util';
+// Increase test timeout since we are fetching the model files from GCS.
+jasmine.DEFAULT_TIMEOUT_INTERVAL = 20000;
+
// Run browser tests againts both the cpu and webgl backends.
setTestEnvs([
// WebGL.
diff --git a/tfjs-automl/src/test_node.ts b/tfjs-automl/src/test_node.ts
index c5dadf988ae..6084df3add7 100644
--- a/tfjs-automl/src/test_node.ts
+++ b/tfjs-automl/src/test_node.ts
@@ -20,6 +20,9 @@ import {setTestEnvs} from '@tensorflow/tfjs-core/dist/jasmine_util';
// tslint:disable-next-line:no-require-imports
const jasmine = require('jasmine');
+// Increase test timeout since we are fetching the model files from GCS.
+jasmine.DEFAULT_TIMEOUT_INTERVAL = 20000;
+
process.on('unhandledRejection', e => {
throw e;
});
diff --git a/tfjs-automl/src/test_util.ts b/tfjs-automl/src/test_util.ts
new file mode 100644
index 00000000000..0a11a0acf63
--- /dev/null
+++ b/tfjs-automl/src/test_util.ts
@@ -0,0 +1,33 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+export async function fetchImage(url: string): Promise {
+ const response = await fetch(url);
+ const blob = await response.blob();
+ const img = new Image();
+ const blobUrl = URL.createObjectURL(blob);
+ return new Promise((resolve, reject) => {
+ img.onload = () => {
+ URL.revokeObjectURL(blobUrl);
+ resolve(img);
+ };
+ img.onerror = (evt /* Arg is an event, not error. Can't rethrow it */) => {
+ reject(new Error('Failed to load blob as image.'));
+ };
+ img.src = blobUrl;
+ });
+}
diff --git a/tfjs-automl/src/types.ts b/tfjs-automl/src/types.ts
index 59b250d2d42..e969b37c063 100644
--- a/tfjs-automl/src/types.ts
+++ b/tfjs-automl/src/types.ts
@@ -17,11 +17,5 @@
import {Tensor3D} from '@tensorflow/tfjs-core';
-export type ClassificationPrediction = Array<{prob: number, label: string}>;
-
-export interface ImageClassificationOptions {
- centerCrop: boolean;
-}
-
export type ImageInput =
ImageData|HTMLImageElement|HTMLCanvasElement|HTMLVideoElement|Tensor3D;
diff --git a/tfjs-automl/src/util.ts b/tfjs-automl/src/util.ts
new file mode 100644
index 00000000000..484e6ad772f
--- /dev/null
+++ b/tfjs-automl/src/util.ts
@@ -0,0 +1,32 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {browser, Tensor, Tensor3D, util} from '@tensorflow/tfjs-core';
+import {ImageInput} from './types';
+
+export function imageToTensor(img: ImageInput): Tensor3D {
+ return img instanceof Tensor ? img : browser.fromPixels(img);
+}
+
+/** Loads and parses the dictionary. */
+export async function loadDictionary(modelUrl: string): Promise {
+ const prefixUrl = modelUrl.slice(0, modelUrl.lastIndexOf('/'));
+ const dictUrl = `${prefixUrl}/dict.txt`;
+ const response = await util.fetch(dictUrl);
+ const text = await response.text();
+ return text.trim().split('\n');
+}
diff --git a/tfjs-automl/yarn.lock b/tfjs-automl/yarn.lock
index c0283681257..c8a2c0a43fd 100644
--- a/tfjs-automl/yarn.lock
+++ b/tfjs-automl/yarn.lock
@@ -268,7 +268,7 @@ atob@^2.1.1:
resolved "https://registry.yarnpkg.com/atob/-/atob-2.1.2.tgz#6d9517eb9e030d2436666651e86bd9f6f13533c9"
integrity sha512-Wm6ukoaOGJi/73p/cl2GvLjTI5JM1k/O14isD73YML8StrH/7/lRFgmg8nICZgD3bZZvjwCGxtMOD3wWNAu8cg==
-babel-code-frame@^6.22.0:
+babel-code-frame@^6.20.0, babel-code-frame@^6.22.0:
version "6.26.0"
resolved "https://registry.yarnpkg.com/babel-code-frame/-/babel-code-frame-6.26.0.tgz#63fd43f7dc1e3bb7ce35947db8fe369a3f58c74b"
integrity sha1-Y/1D99weO7fONZR9uP42mj9Yx0s=
@@ -707,7 +707,7 @@ color-name@1.1.3:
resolved "https://registry.yarnpkg.com/color-name/-/color-name-1.1.3.tgz#a7d0558bd89c42f795dd42328f740831ca53bc25"
integrity sha1-p9BVi9icQveV3UIyj3QIMcpTvCU=
-colors@^1.1.0:
+colors@^1.1.0, colors@^1.1.2:
version "1.3.3"
resolved "https://registry.yarnpkg.com/colors/-/colors-1.3.3.tgz#39e005d546afe01e01f9c4ca8fa50f686a01205d"
integrity sha512-mmGt/1pZqYRjMxB1axhTo16/snVZ5krrKkcmMeVKxzECMMXoCgnvTPp10QgHfcbQZw8Dq2jMNG6je4JlWU0gWg==
@@ -1011,7 +1011,7 @@ di@^0.0.1:
resolved "https://registry.yarnpkg.com/di/-/di-0.0.1.tgz#806649326ceaa7caa3306d75d985ea2748ba913c"
integrity sha1-gGZJMmzqp8qjMG112YXqJ0i6kTw=
-diff@^3.2.0:
+diff@^3.0.1, diff@^3.2.0:
version "3.5.0"
resolved "https://registry.yarnpkg.com/diff/-/diff-3.5.0.tgz#800c0dd1e0a8bfbc95835c202ad220fe317e5a12"
integrity sha512-A46qtFgd+g7pDZinpnwiRJtxbC1hpgf0uzP3iG89scHk0AUC7A1TGxf5OiiOUv/JMZR8GOt8hL900hV0bOy5xA==
@@ -1318,6 +1318,13 @@ find-up@^1.0.0:
path-exists "^2.0.0"
pinkie-promise "^2.0.0"
+findup-sync@~0.3.0:
+ version "0.3.0"
+ resolved "https://registry.yarnpkg.com/findup-sync/-/findup-sync-0.3.0.tgz#37930aa5d816b777c03445e1966cc6790a4c0b16"
+ integrity sha1-N5MKpdgWt3fANEXhlmzGeQpMCxY=
+ dependencies:
+ glob "~5.0.0"
+
flatted@^2.0.0:
version "2.0.1"
resolved "https://registry.yarnpkg.com/flatted/-/flatted-2.0.1.tgz#69e57caa8f0eacbc281d2e2cb458d46fdb449e08"
@@ -1448,7 +1455,7 @@ glob-parent@^3.1.0:
is-glob "^3.1.0"
path-dirname "^1.0.0"
-glob@^5.0.15:
+glob@^5.0.15, glob@~5.0.0:
version "5.0.15"
resolved "https://registry.yarnpkg.com/glob/-/glob-5.0.15.tgz#1bc936b9e02f4a603fcc222ecf7633d30b8b93b1"
integrity sha1-G8k2ueAvSmA/zCIuz3Yz0wuLk7E=
@@ -2575,7 +2582,7 @@ once@1.x, once@^1.3.0:
dependencies:
wrappy "1"
-optimist@^0.6.1:
+optimist@^0.6.1, optimist@~0.6.0:
version "0.6.1"
resolved "https://registry.yarnpkg.com/optimist/-/optimist-0.6.1.tgz#da3ea74686fa21a19a111c326e90eb15a0196686"
integrity sha1-2j6nRob6IaGaERwybpDrFaAZZoY=
@@ -2988,7 +2995,7 @@ resolve@1.1.7, resolve@1.1.x:
resolved "https://registry.yarnpkg.com/resolve/-/resolve-1.1.7.tgz#203114d82ad2c5ed9e8e0411b3932875e889e97b"
integrity sha1-IDEU2CrSxe2ejgQRs5ModeiJ6Xs=
-resolve@^1.1.6, resolve@^1.10.0, resolve@^1.3.2, resolve@^1.5.0, resolve@^1.7.1:
+resolve@^1.1.6, resolve@^1.1.7, resolve@^1.10.0, resolve@^1.3.2, resolve@^1.5.0, resolve@^1.7.1:
version "1.12.0"
resolved "https://registry.yarnpkg.com/resolve/-/resolve-1.12.0.tgz#3fc644a35c84a48554609ff26ec52b66fa577df6"
integrity sha512-B/dOmuoAik5bKcD6s6nXDCjzUKnaDvdkRyAk6rsmsKLipWj4797iothd7jmmUhWTfinVMU+wc56rYKsit2Qy4w==
@@ -3557,6 +3564,20 @@ trim-newlines@^1.0.0:
resolved "https://registry.yarnpkg.com/trim-newlines/-/trim-newlines-1.0.0.tgz#5887966bb582a4503a41eb524f7d35011815a613"
integrity sha1-WIeWa7WCpFA6QetST301ARgVphM=
+ts-lint@^4.5.1:
+ version "4.5.1"
+ resolved "https://registry.yarnpkg.com/ts-lint/-/ts-lint-4.5.1.tgz#9c22b7b7b862b67324dd1bd213a845c03a7fb8c0"
+ integrity sha1-nCK3t7hitnMk3RvSE6hFwDp/uMA=
+ dependencies:
+ babel-code-frame "^6.20.0"
+ colors "^1.1.2"
+ diff "^3.0.1"
+ findup-sync "~0.3.0"
+ glob "^7.1.1"
+ optimist "~0.6.0"
+ resolve "^1.1.7"
+ tsutils "^1.1.0"
+
ts-node@^8.3.0:
version "8.3.0"
resolved "https://registry.yarnpkg.com/ts-node/-/ts-node-8.3.0.tgz#e4059618411371924a1fb5f3b125915f324efb57"
@@ -3596,6 +3617,11 @@ tslint@~5.11.0:
tslib "^1.8.0"
tsutils "^2.27.2"
+tsutils@^1.1.0:
+ version "1.9.1"
+ resolved "https://registry.yarnpkg.com/tsutils/-/tsutils-1.9.1.tgz#b9f9ab44e55af9681831d5f28d0aeeaf5c750cb0"
+ integrity sha1-ufmrROVa+WgYMdXyjQrur1x1DLA=
+
tsutils@^2.27.2:
version "2.29.0"
resolved "https://registry.yarnpkg.com/tsutils/-/tsutils-2.29.0.tgz#32b488501467acbedd4b85498673a0812aca0b99"
+
+