[tfjs-react-native] webcam demo (#1938)

INTERNAL

Author: tafsiri

tfjs-react-native/integration_rn59/App.tsx

@@ -16,18 +16,19 @@
  */
 
 import React, { Fragment } from 'react';
-import { Button, SafeAreaView, StyleSheet, ScrollView, View, Text, StatusBar } from 'react-native';
+import { Button, SafeAreaView, StyleSheet, View, Text, StatusBar } from 'react-native';
 
 import * as tf from '@tensorflow/tfjs';
 import '@tensorflow/tfjs-react-native';
 
 import { Diagnostic } from './components/diagnostic';
 import { MobilenetDemo } from './components/mobilenet_demo';
 import { TestRunner } from './components/tfjs_unit_test_runner';
+import { WebcamDemo } from './components/webcam/webcam_demo';
 
-const BACKEND_TO_USE = 'cpu';
+const BACKEND_TO_USE = 'rn-webgl';
 
-export type Screen = 'main' | 'diag' | 'demo' | 'test';
+export type Screen = 'main' | 'diag' | 'demo' | 'test' | 'webcam';
 
 interface AppState {
   isTfReady: boolean;
@@ -46,6 +47,7 @@ export class App extends React.Component<{}, AppState> {
     this.showDemoScreen = this.showDemoScreen.bind(this);
     this.showMainScreen = this.showMainScreen.bind(this);
     this.showTestScreen = this.showTestScreen.bind(this);
+    this.showWebcamDemo= this.showWebcamDemo.bind(this);
   }
 
   async componentDidMount() {
@@ -72,6 +74,10 @@ export class App extends React.Component<{}, AppState> {
     this.setState({ currentScreen: 'test' });
   }
 
+  showWebcamDemo() {
+    this.setState({ currentScreen: 'webcam' });
+  }
+
   renderMainScreen() {
     return <Fragment>
       <View style={styles.sectionContainer}>
@@ -95,6 +101,13 @@ export class App extends React.Component<{}, AppState> {
           title='Show Test Screen'
         />
       </View>
+      <View style={styles.sectionContainer}>
+        <Text style={styles.sectionTitle}>Webcam Demo</Text>
+        <Button
+          onPress={this.showWebcamDemo}
+          title='Show Webcam Demo'
+        />
+      </View>
     </Fragment>;
   }
 
@@ -119,6 +132,12 @@ export class App extends React.Component<{}, AppState> {
     </Fragment>;
   }
 
+  renderWebcamDemo() {
+    return <Fragment>
+      <WebcamDemo returnToMain={this.showMainScreen}/>
+    </Fragment>;
+  }
+
   renderLoadingTF() {
     return <Fragment>
       <View style={styles.sectionContainer}>
@@ -139,6 +158,8 @@ export class App extends React.Component<{}, AppState> {
           return this.renderDemoScreen();
         case 'test':
           return this.renderTestScreen();
+        case 'webcam':
+          return this.renderWebcamDemo();
         default:
           return this.renderMainScreen();
       }
@@ -153,13 +174,9 @@ export class App extends React.Component<{}, AppState> {
       <Fragment>
         <StatusBar barStyle='dark-content' />
         <SafeAreaView>
-          <ScrollView
-            contentInsetAdjustmentBehavior='automatic'
-            style={styles.scrollView}>
-            <View style={styles.body}>
-              {this.renderContent()}
-            </View>
-          </ScrollView>
+          <View style={styles.body}>
+            {this.renderContent()}
+          </View>
         </SafeAreaView>
       </Fragment>
     );

tfjs-react-native/integration_rn59/android/app/src/main/java/com/integration_rn59/generated/BasePackageList.java

@@ -8,8 +8,10 @@ public class BasePackageList {
   public List<Package> getPackageList() {
     return Arrays.<Package>asList(
         new expo.modules.gl.GLPackage(),
+        new expo.modules.camera.CameraPackage(),
         new expo.modules.constants.ConstantsPackage(),
         new expo.modules.filesystem.FileSystemPackage(),
+        new expo.modules.imagemanipulator.ImageManipulatorPackage(),
         new expo.modules.permissions.PermissionsPackage()
     );
   }

tfjs-react-native/integration_rn59/android/build.gradle

@@ -45,6 +45,12 @@ allprojects {
         maven {
             // All of React Native (JS, Obj-C sources, Android binaries) is installed from npm
             url "$rootDir/../node_modules/react-native/android"
+
+        }
+        maven {
+          // expo-camera bundles a custom com.google.android:cameraview
+            url "$rootDir/../node_modules/expo-camera/android/maven"
         }
+
     }
 }

tfjs-react-native/integration_rn59/components/webcam/image_utils.ts

@@ -0,0 +1,85 @@
+/**
+ * @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 * as tf from '@tensorflow/tfjs';
+import * as ImageManipulator from 'expo-image-manipulator';
+import * as jpeg from 'jpeg-js';
+
+export function toDataUri(base64: string): string {
+  return `data:image/jpeg;base64,${base64}`;
+}
+
+export async function resizeImage(
+    imageUrl: string, width: number): Promise<ImageManipulator.ImageResult> {
+  const actions = [{
+    resize: {
+      width,
+    },
+  }];
+  const saveOptions = {
+    compress: 0.75,
+    format: ImageManipulator.SaveFormat.JPEG,
+    base64: true,
+  };
+  const res =
+      await ImageManipulator.manipulateAsync(imageUrl, actions, saveOptions);
+  return res;
+}
+
+export async function base64ImageToTensor(base64: string):
+    Promise<tf.Tensor3D> {
+  const rawImageData = tf.util.encodeString(base64, 'base64');
+  const TO_UINT8ARRAY = true;
+  const {width, height, data} = jpeg.decode(rawImageData, TO_UINT8ARRAY);
+  // Drop the alpha channel info
+  const buffer = new Uint8Array(width * height * 3);
+  let offset = 0;  // offset into original data
+  for (let i = 0; i < buffer.length; i += 3) {
+    buffer[i] = data[offset];
+    buffer[i + 1] = data[offset + 1];
+    buffer[i + 2] = data[offset + 2];
+
+    offset += 4;
+  }
+  return tf.tensor3d(buffer, [height, width, 3]);
+}
+
+export async function tensorToImageUrl(imageTensor: tf.Tensor3D):
+    Promise<string> {
+  const [height, width] = imageTensor.shape;
+  const buffer = await imageTensor.toInt().data();
+  const frameData = new Uint8Array(width * height * 4);
+
+  let offset = 0;
+  for (let i = 0; i < frameData.length; i += 4) {
+    frameData[i] = buffer[offset];
+    frameData[i + 1] = buffer[offset + 1];
+    frameData[i + 2] = buffer[offset + 2];
+    frameData[i + 3] = 0xFF;
+
+    offset += 3;
+  }
+
+  const rawImageData = {
+    data: frameData,
+    width,
+    height,
+  };
+  const jpegImageData = jpeg.encode(rawImageData, 75);
+  const base64Encoding = tf.util.decodeString(jpegImageData.data, 'base64');
+  return base64Encoding;
+}

tfjs-react-native/integration_rn59/components/webcam/style_transfer.ts

@@ -0,0 +1,105 @@
+/**
+ * @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 * as tf from '@tensorflow/tfjs';
+
+const STYLENET_URL =
+    'https://cdn.jsdelivr.net/gh/reiinakano/arbitrary-image-stylization-tfjs@master/saved_model_style_js/model.json';
+const TRANSFORMNET_URL =
+    'https://cdn.jsdelivr.net/gh/reiinakano/arbitrary-image-stylization-tfjs@master/saved_model_transformer_separable_js/model.json';
+
+export class StyleTranfer {
+  private styleNet?: tf.GraphModel;
+  private transformNet?: tf.GraphModel;
+
+  constructor() {}
+
+  async init() {
+    await Promise.all([this.loadStyleModel(), this.loadTransformerModel()]);
+    await this.warmup();
+  }
+
+  async loadStyleModel() {
+    if (this.styleNet == null) {
+      this.styleNet = await tf.loadGraphModel(STYLENET_URL);
+      console.log('stylenet loaded');
+    }
+  }
+
+  async loadTransformerModel() {
+    if (this.transformNet == null) {
+      this.transformNet = await tf.loadGraphModel(TRANSFORMNET_URL);
+      console.log('transformnet loaded');
+    }
+  }
+
+  async warmup() {
+    // Also warmup
+    const input: tf.Tensor3D = tf.randomNormal([320, 240, 3]);
+    const res = this.stylize(input, input);
+    await res.data();
+    tf.dispose([input, res]);
+  }
+
+  /**
+   * This function returns style bottleneck features for
+   * the given image.
+   *
+   * @param style Style image to get 100D bottleneck features for
+   */
+  private predictStyleParameters(styleImage: tf.Tensor3D): tf.Tensor4D {
+    return tf.tidy(() => {
+      if (this.styleNet == null) {
+        throw new Error('Stylenet not loaded');
+      }
+      return this.styleNet.predict(
+          styleImage.toFloat().div(tf.scalar(255)).expandDims());
+    }) as tf.Tensor4D;
+  }
+
+  /**
+   * This function stylizes the content image given the bottleneck
+   * features. It returns a tf.Tensor3D containing the stylized image.
+   *
+   * @param content Content image to stylize
+   * @param bottleneck Bottleneck features for the style to use
+   */
+  private produceStylized(contentImage: tf.Tensor3D, bottleneck: tf.Tensor4D):
+      tf.Tensor3D {
+    return tf.tidy(() => {
+      if (this.transformNet == null) {
+        throw new Error('Transformnet not loaded');
+      }
+      const input = contentImage.toFloat().div(tf.scalar(255)).expandDims();
+      const image: tf.Tensor4D =
+          this.transformNet.predict([input, bottleneck]) as tf.Tensor4D;
+      return image.mul(255).squeeze();
+    });
+  }
+
+  public stylize(styleImage: tf.Tensor3D, contentImage: tf.Tensor3D):
+      tf.Tensor3D {
+    const start = Date.now();
+    console.log(styleImage.shape, contentImage.shape);
+    const styleRepresentation = this.predictStyleParameters(styleImage);
+    const stylized = this.produceStylized(contentImage, styleRepresentation);
+    tf.dispose([styleRepresentation]);
+    const end = Date.now();
+    console.log('stylization scheduled', end - start);
+    return stylized;
+  }
+}