Move minDepth and maxDepth to estimation config (#1014)

Author: ahmedsabie

depth-estimation/README.md

@@ -33,17 +33,17 @@ For example:
 
 ```javascript
 const model = depthEstimation.SupportedModels.ARPortraitDepth;
-const estimatorConfig = {
-  minDepth: 0,
-  maxDepth: 1,
-}
-const estimator = await depthEstimation.createEstimator(model, estimatorConfig);
+const estimator = await depthEstimation.createEstimator(model);
 ```
 
 Next, you can use the estimator to estimate depth.
 
 ```javascript
-const depthMap = await estimator.estimateDepth(image);
+const estimationConfig = {
+  minDepth: 0,
+  maxDepth: 1,
+}
+const depthMap = await estimator.estimateDepth(image, estimationConfig);
 ```
 
 The returned depth map contains depth values for each pixel in the image.

depth-estimation/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tensorflow-models/depth-estimation",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "description": "Pretrained depth model",
   "main": "dist/index.js",
   "jsnext:main": "dist/depth-estimation.esm.js",

depth-estimation/src/ar_portrait_depth/README.md

@@ -62,12 +62,6 @@ Pass in `depthEstimation.SupportedModels.ARPortraitDepth` from the
 
 `estimatorConfig` is an object that defines ARPortraitDepth specific configurations for `ARPortraitDepthModelConfig`:
 
-*   *minDepth*: The minimum depth value for the model to map to 0. Any smaller
-depth values will also get mapped to 0.
-
-*   *maxDepth*: The maximum depth value for the model to map to 1. Any larger
-depth values will also get mapped to 1.
-
 *   *segmentationModelUrl*: An optional string that specifies custom url of
 the segmenter model. This is useful for area/countries that don't have access to the model hosted on tf.hub. It also accepts `io.IOHandler` which can be used with
 [tfjs-react-native](https://github.com/tensorflow/tfjs/tree/master/tfjs-react-native)
@@ -97,6 +91,12 @@ options, you can pass in a second `estimationConfig` parameter.
 
 `estimationConfig` is an object that defines ARPortraitDepth specific configurations for `ARPortraitDepthEstimationConfig`:
 
+*   *minDepth*: The minimum depth value for the model to map to 0. Any smaller
+depth values will also get mapped to 0.
+
+*   *maxDepth*: The maximum depth value for the model to map to 1. Any larger
+depth values will also get mapped to 1.
+
 *   *flipHorizontal*: Optional. Defaults to false. When image data comes from camera, the result has to flip horizontally.
 
 The following code snippet demonstrates how to run the model inference:

depth-estimation/src/ar_portrait_depth/ar_portrait_depth_test.ts

@@ -114,13 +114,13 @@ describeWithFlags('ARPortraitDepth', ALL_ENVS, () => {
 
     // Note: this makes a network request for model assets.
     const estimator = await depthEstimation.createEstimator(
-        depthEstimation.SupportedModels.ARPortraitDepth,
-        {minDepth: 0, maxDepth: 1});
+        depthEstimation.SupportedModels.ARPortraitDepth);
     const input: tf.Tensor3D = tf.zeros([128, 128, 3]);
 
     const beforeTensors = tf.memory().numTensors;
 
-    const depthMap = await estimator.estimateDepth(input);
+    const depthMap =
+        await estimator.estimateDepth(input, {minDepth: 0, maxDepth: 1});
 
     (await depthMap.toTensor()).dispose();
     expect(tf.memory().numTensors).toEqual(beforeTensors);
@@ -133,21 +133,25 @@ describeWithFlags('ARPortraitDepth', ALL_ENVS, () => {
 
   it('throws error when minDepth is not set.', async (done) => {
     try {
-      await depthEstimation.createEstimator(
+      const estimator = await depthEstimation.createEstimator(
           depthEstimation.SupportedModels.ARPortraitDepth);
+      const input: tf.Tensor3D = tf.zeros([128, 128, 3]);
+      await estimator.estimateDepth(input);
       done.fail('Loading without minDepth succeeded unexpectedly.');
     } catch (e) {
       expect(e.message).toEqual(
-          'A model config with minDepth and maxDepth set must be provided.');
+          'An estimation config with ' +
+          'minDepth and maxDepth set must be provided.');
       done();
     }
   });
 
   it('throws error when minDepth is greater than maxDepth.', async (done) => {
     try {
-      await depthEstimation.createEstimator(
-          depthEstimation.SupportedModels.ARPortraitDepth,
-          {minDepth: 1, maxDepth: 0.99});
+      const estimator = await depthEstimation.createEstimator(
+          depthEstimation.SupportedModels.ARPortraitDepth);
+      const input: tf.Tensor3D = tf.zeros([128, 128, 3]);
+      await estimator.estimateDepth(input, {minDepth: 1, maxDepth: 0.99});
       done.fail(
           'Loading with minDepth greater than maxDepth ' +
           'succeeded unexpectedly.');
@@ -187,12 +191,12 @@ describeWithFlags('ARPortraitDepth static image ', BROWSER_ENVS, () => {
     // Get actual depth values.
     // Note: this makes a network request for model assets.
     estimator = await depthEstimation.createEstimator(
-        depthEstimation.SupportedModels.ARPortraitDepth,
-        {minDepth: 0.2, maxDepth: 0.9});
+        depthEstimation.SupportedModels.ARPortraitDepth);
 
     const beforeTensors = tf.memory().numTensors;
 
-    const result = await estimator.estimateDepth(image);
+    const result =
+        await estimator.estimateDepth(image, {minDepth: 0.2, maxDepth: 0.9});
     const actualDepthValues = await result.toTensor();
     const coloredDepthValues =
         actualDepthValues.arraySync().flat().map(value => turboPlus(value));

depth-estimation/src/ar_portrait_depth/constants.ts

@@ -14,12 +14,16 @@
  * limitations under the License.
  * =============================================================================
  */
-import {ARPortraitDepthEstimationConfig} from './types';
+import {ARPortraitDepthModelConfig} from './types';
 
 export const DEFAULT_AR_PORTRAIT_DEPTH_MODEL_URL =
     'https://tfhub.dev/tensorflow/tfjs-model/ar_portrait_depth/1';
 
-export const DEFAULT_AR_PORTRAIT_DEPTH_ESTIMATION_CONFIG:
-    ARPortraitDepthEstimationConfig = {
-      flipHorizontal: false,
+export const DEFAULT_AR_PORTRAIT_DEPTH_MODEL_CONFIG:
+    ARPortraitDepthModelConfig = {
+      depthModelUrl: DEFAULT_AR_PORTRAIT_DEPTH_MODEL_URL,
     };
+
+export const DEFAULT_AR_PORTRAIT_DEPTH_ESTIMATION_CONFIG = {
+  flipHorizontal: false,
+};

depth-estimation/src/ar_portrait_depth/estimator.ts

@@ -56,8 +56,7 @@ const PORTRAIT_WIDTH = 192;
 class ARPortraitDepthEstimator implements DepthEstimator {
   constructor(
       private readonly segmenter: bodySegmentation.BodySegmenter,
-      private readonly estimatorModel: tfconv.GraphModel,
-      private readonly minDepth: number, private readonly maxDepth: number) {}
+      private readonly estimatorModel: tfconv.GraphModel) {}
 
   /**
    * Estimates depth for an image or video frame.
@@ -69,8 +68,14 @@ class ARPortraitDepthEstimator implements DepthEstimator {
    * image to feed through the network.
    *
    * @param config Optional.
+   *       minDepth: The minimum depth value for the model to map to 0. Any
+   * smaller depth values will also get mapped to 0.
+   *
+   *       maxDepth`: The maximum depth value for the model to map to 1. Any
+   * larger depth values will also get mapped to 1.
+   *
    *       flipHorizontal: Optional. Default to false. When image data comes
-   *       from camera, the result has to flip horizontally.
+   * from camera, the result has to flip horizontally.
    *
    * @return `DepthMap`.
    */
@@ -128,7 +133,7 @@ class ARPortraitDepthEstimator implements DepthEstimator {
 
       // Normalize to user requirements.
       const depthTransform =
-          transformValueRange(this.minDepth, this.maxDepth, 0, 1);
+          transformValueRange(config.minDepth, config.maxDepth, 0, 1);
 
       // depth4D is roughly in [0,2] range, so half the scale factor to put it
       // in [0,1] range.
@@ -188,6 +193,5 @@ export async function load(modelConfig: ARPortraitDepthModelConfig):
       bodySegmentation.SupportedModels.MediaPipeSelfieSegmentation,
       {runtime: 'tfjs', modelUrl: config.segmentationModelUrl});
 
-  return new ARPortraitDepthEstimator(
-      segmenter, depthModel, config.minDepth, config.maxDepth);
+  return new ARPortraitDepthEstimator(segmenter, depthModel);
 }

depth-estimation/src/ar_portrait_depth/estimator_utils.ts

@@ -15,34 +15,35 @@
  * =============================================================================
  */
 
-import {DEFAULT_AR_PORTRAIT_DEPTH_ESTIMATION_CONFIG, DEFAULT_AR_PORTRAIT_DEPTH_MODEL_URL} from './constants';
+import {DEFAULT_AR_PORTRAIT_DEPTH_ESTIMATION_CONFIG, DEFAULT_AR_PORTRAIT_DEPTH_MODEL_CONFIG} from './constants';
 import {ARPortraitDepthEstimationConfig, ARPortraitDepthModelConfig} from './types';
 
 export function validateModelConfig(modelConfig: ARPortraitDepthModelConfig):
     ARPortraitDepthModelConfig {
-  if (modelConfig == null || modelConfig.minDepth == null ||
-      modelConfig.maxDepth == null) {
-    throw new Error(
-        `A model config with minDepth and maxDepth set must be provided.`);
-  }
-
-  if (modelConfig.minDepth > modelConfig.maxDepth) {
-    throw new Error('minDepth must be <= maxDepth.');
+  if (modelConfig == null) {
+    return {...DEFAULT_AR_PORTRAIT_DEPTH_MODEL_CONFIG};
   }
 
   const config = {...modelConfig};
 
   if (config.depthModelUrl == null) {
-    config.depthModelUrl = DEFAULT_AR_PORTRAIT_DEPTH_MODEL_URL;
+    config.depthModelUrl = DEFAULT_AR_PORTRAIT_DEPTH_MODEL_CONFIG.depthModelUrl;
   }
 
   return config;
 }
 
 export function validateEstimationConfig(
     estimationConfig: ARPortraitDepthEstimationConfig) {
-  if (estimationConfig == null) {
-    return {...DEFAULT_AR_PORTRAIT_DEPTH_ESTIMATION_CONFIG};
+  if (estimationConfig == null || estimationConfig.minDepth == null ||
+      estimationConfig.maxDepth == null) {
+    throw new Error(
+        'An estimation config with ' +
+        'minDepth and maxDepth set must be provided.');
+  }
+
+  if (estimationConfig.minDepth > estimationConfig.maxDepth) {
+    throw new Error('minDepth must be <= maxDepth.');
   }
 
   const config = {...estimationConfig};

depth-estimation/src/ar_portrait_depth/types.ts

@@ -22,12 +22,6 @@ import {EstimationConfig, ModelConfig} from '../types';
 /**
  * Model parameters for ARPortraitDepth.
  *
- * `minDepth`: The minimum depth value for the model to map to 0. Any smaller
- * depth values will also get mapped to 0.
- *
- * `maxDepth`: The maximum depth value for the model to map to 1. Any larger
- * depth values will also get mapped to 1.
- *
  * `segmentationModelUrl`: Optional. An optional string that specifies custom
  * url of the selfie segmentation model. This is useful for area/countries that
  * don't have access to the model hosted on tf.hub.

depth-estimation/src/types.ts

@@ -24,26 +24,25 @@ export enum SupportedModels {
 
 /**
  * Common config to create the depth estimator.
+ */
+export interface ModelConfig {}
+
+/**
+ * Common config for the `estimateDepth` method.
  *
  * `minDepth`: The minimum depth value for the model to map to 0. Any smaller
  * depth values will also get mapped to 0.
  *
  * `maxDepth`: The maximum depth value for the model to map to 1. Any larger
  * depth values will also get mapped to 1.
- */
-export interface ModelConfig {
-  minDepth: number;
-  maxDepth: number;
-}
-
-/**
- * Common config for the `estimateDepth` method.
  *
  * `flipHorizontal`: Optional. Default to false. In some cases, the image is
  * mirrored, e.g., video stream from camera, flipHorizontal will flip the
  * keypoints horizontally.
  */
 export interface EstimationConfig {
+  minDepth: number;
+  maxDepth: number;
   flipHorizontal?: boolean;
 }
 

depth-estimation/src/version.ts

@@ -1,5 +1,5 @@
 /** @license See the LICENSE file. */
 
 // This code is auto-generated, do not modify this file!
-const version = '0.0.1';
+const version = '0.0.2';
 export {version};