tfjs-react-native: Memory problem while running the webcam demo

[devops35]

I can't run the webcam demo on my real phone (iphone 6). It exceeds 650 mb ram. Then my application crashes. I can run it in the simulator, it works very, very slowly there. It takes minutes to work. In addition, the ram used is about 2 GB.

my package.json;

"@react-native-async-storage/async-storage": "^1.13.2",
    "@react-native-community/async-storage": "^1.12.1",
    "@tensorflow-models/blazeface": "^0.0.5",
    "@tensorflow-models/mobilenet": "^2.0.4",
    "@tensorflow-models/posenet": "^2.2.1",
    "@tensorflow/tfjs": "^2.8.0",
    "@tensorflow/tfjs-react-native": "^0.5.0",
    "@unimodules/core": "^6.0.0",
    "expo-asset": "^8.2.1",
    "expo-camera": "^9.1.1",
    "expo-constants": "^9.3.5",
    "expo-gl": "^9.2.0",
    "expo-image-manipulator": "^8.4.0",
    "expo-image-picker": "^9.2.1",
    "expo-media-library": "^10.0.0",
    "expo-permissions": "^10.0.0",
    "jpeg-js": "^0.4.2",
    "react": "16.13.1",
    "react-native": "0.63.4",
    "react-native-fs": "^2.16.6",
    "react-native-svg": "^12.1.0",
    "react-native-unimodules": "^0.12.0"

webcam_demo.tsx;

import React from 'react';
import { ActivityIndicator, StyleSheet, View, Image, Text, TouchableHighlight } from 'react-native';
import * as Permissions from 'expo-permissions';
import { Camera } from 'expo-camera';
import {StyleTranfer} from './style_transfer';
import {base64ImageToTensor, tensorToImageUrl, resizeImage, toDataUri} from './image_utils';
import * as tf from '@tensorflow/tfjs';
import * as jpeg from 'jpeg-js'
import { fetch } from '@tensorflow/tfjs-react-native'

interface ScreenProps {
  returnToMain: () => void;
}

interface ScreenState {
  mode: 'results' | 'newStyleImage' | 'newContentImage';
  resultImage?: string;
  styleImage?: string;
  contentImage?: string;
  styleImageUri?: string;
  contentImageUri?: string;
  hasCameraPermission?: boolean;
  // tslint:disable-next-line: no-any
  cameraType: any;
  isLoading: boolean;
}

export default class WebcamDemo extends React.Component<ScreenProps,ScreenState> {
  private camera?: Camera|null;
  private styler: StyleTranfer;

  constructor(props: ScreenProps) {
    super(props);
    this.state = {
      mode: 'results',
      cameraType: Camera.Constants.Type.back,
      isLoading: true,
    };
    this.styler = new StyleTranfer();
  }

  async componentDidMount() {
    // await tf.setBackend('cpu')
    // await tf.setBackend('rn-webgl')
    await tf.ready()
    // tf.enableDebugMode()
    // await tf.profile();

    await this.styler.init();
    const { status } = await Permissions.askAsync(Permissions.CAMERA);
    this.setState({
      hasCameraPermission: status === 'granted',
      isLoading: false
    });
  }

  showResults() {
    this.setState({ mode: 'results' });
  }

  takeStyleImage() {
    this.setState({ mode: 'newStyleImage' });
  }

  takeContentImage() {
    this.setState({ mode: 'newContentImage' });
  }

  flipCamera() {
    const newState = this.state.cameraType === Camera.Constants.Type.back
          ? Camera.Constants.Type.front
          : Camera.Constants.Type.back;
    this.setState({
      cameraType: newState,
    });
  }

  renderStyleImagePreview() {
    const {styleImage} = this.state;
    if(styleImage == null) {
      return (
        <View>
          <Text style={styles.instructionText}>Style</Text>
          <Text style={{fontSize: 48, paddingLeft: 0}}>💅🏽</Text>
        </View>
      );
    } else {
      return (
        <View>
          <Image
            style={styles.imagePreview}
            source={{uri: toDataUri(styleImage)}} />
            <Text style={styles.centeredText}>Style</Text>
        </View>
      );
    }
  }

  renderContentImagePreview() {
    const {contentImage} = this.state;
    if(contentImage == null) {
      return (
        <View>
          <Text style={styles.instructionText}>Stuff</Text>
          <Text style={{fontSize: 48, paddingLeft: 0}}>🖼️</Text>
        </View>
      );
    } else {
      return (
        <View>
          <Image
            style={styles.imagePreview}
            source={{uri: toDataUri(contentImage)}} />
            <Text style={styles.centeredText}>Stuff</Text>
        </View>
      );
    }
  }

  async imageToTensor(rawImageData1) {
    // const imageAssetPath = Image.resolveAssetSource({uri:rawImageData1})
    const imageAssetPath = Image.resolveAssetSource(rawImageData1)

    console.log("imageAssetPath",imageAssetPath)

    const response = await fetch(imageAssetPath.uri, {}, { isBinary: true })
    const rawImageData = await response.arrayBuffer()
    const TO_UINT8ARRAY = true
    const { width, height, data } = jpeg.decode(rawImageData, TO_UINT8ARRAY)
    // Drop the alpha channel info for mobilenet
    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])
  }

  async stylize(contentImage, styleImage){

    // console.log("stylize1",this.state.contentImageUri, contentImage)
    console.log("aaaa")

    // const contentTensor = await this.imageToTensor(contentImage);
    const contentTensor = await this.imageToTensor(require("../ts_models/images/statue_of_liberty.jpg"));
    console.log("bbbb")
    // const styleTensor = await this.imageToTensor(styleImage);
    const styleTensor = await this.imageToTensor(require("../ts_models/images/udnie.jpg"));
    console.log("cccc")
    const stylizedResult = this.styler.stylize(
      styleTensor, contentTensor);
      console.log("dddd")
    const stylizedImage = await tensorToImageUrl(stylizedResult); 
    console.log("eeee")
    tf.dispose([contentTensor, styleTensor, stylizedResult]);
    console.log("ffff")
    return stylizedImage;
  }

  async handleCameraCapture() {
    const {mode} = this.state;
    let {styleImage, contentImage, resultImage, contentImageUri, styleImageUri} = this.state;
    this.setState({
      isLoading: true,
    });
    let image = await this.camera!.takePictureAsync({
      skipProcessing: true,
    });
    image = await resizeImage(image.uri, 240);

    // console.log("image",image)

    if(mode === 'newStyleImage' && image.base64 != null) {
      styleImage = image.base64;
      styleImageUri = image.uri;
      if(contentImage == null) {

        console.log("image 1111",image.uri)

        this.setState({
          styleImage,
          styleImageUri:image.uri,
          mode: 'results',
          isLoading: false,
        });
      } else {
        console.log("image 2222",image.uri)
        resultImage = await this.stylize(contentImageUri, styleImageUri),
        this.setState({
          styleImage,
          contentImage,
          styleImageUri,
          contentImageUri,
        //   contentImageUri:image.uri,
          resultImage,
          mode: 'results',
          isLoading: false,
        });
      }
    } else if (mode === 'newContentImage' && image.base64 != null) {
        console.log("image 3333",image.uri)
      contentImage = image.base64;
      contentImageUri = image.uri;
      if(styleImage == null) {
        console.log("image 5555",image.uri)
        this.setState({
          contentImage,
          contentImageUri:image.uri,
          mode: 'results',
          isLoading: false,
        });
      } else {
        console.log("image 4444",contentImageUri,styleImageUri)
        resultImage = await this.stylize(contentImageUri, styleImageUri);
        this.setState({
          contentImage,
          styleImage,
          styleImageUri,
          contentImageUri,
          resultImage,
          mode: 'results',
          isLoading: false,
        });
      }
    }
  }

  renderCameraCapture() {
    const {hasCameraPermission} = this.state;

    if (hasCameraPermission === null) {
      return <View />;
    } else if (hasCameraPermission === false) {
      return <Text>No access to camera</Text>;
    }
    return (
      <View  style={styles.cameraContainer}>
        <Camera
          style={styles.camera}
          type={this.state.cameraType}
          ref={ref => { this.camera = ref; }}>
        </Camera>
        <View style={styles.cameraControls}>
            <TouchableHighlight
              style={styles.flipCameraBtn}
              onPress={() => {this.flipCamera();}}
              underlayColor='#FFDE03'>
              <Text style={{fontSize: 16, color: 'white'}}>
                FLIP
              </Text>
            </TouchableHighlight>
            <TouchableHighlight
              style={styles.takeImageBtn}
              onPress={() => { this.handleCameraCapture(); }}
              underlayColor='#FFDE03'>
              <Text style={{fontSize: 16, color: 'white', fontWeight: 'bold'}}>
                TAKE
              </Text>
            </TouchableHighlight>
            <TouchableHighlight
              style={styles.cancelBtn}
              onPress={() => {this.showResults(); }}
              underlayColor='#FFDE03'>
              <Text style={{fontSize: 16, color: 'white'}}>
                BACK
              </Text>
            </TouchableHighlight>
          </View>
        </View>
    );
  }

  renderResults() {
    const {resultImage} = this.state;
    return (
      <View>
        <View style={styles.resultImageContainer}>
          {resultImage == null ?
            <Text style={styles.introText}>
              Tap the squares below to add style and content
              images and see the magic!
            </Text>
            :
            <Image
              style={styles.resultImage}
              resizeMode='contain'
              source={{uri: toDataUri(resultImage)}} />
          }
          <TouchableHighlight
            style={styles.styleImageContainer}
            onPress={() => this.takeStyleImage()}
            underlayColor='white'>
              {this.renderStyleImagePreview()}
          </TouchableHighlight>

          <TouchableHighlight
            style={styles.contentImageContainer}
            onPress={() => this.takeContentImage()}
            underlayColor='white'>
            {this.renderContentImagePreview()}
          </TouchableHighlight>

        </View>
      </View>
    );
  }

  render() {
    const {mode, isLoading} = this.state;
    return (
      <View style={{width:'100%'}}>
        {isLoading ? <View style={[styles.loadingIndicator]}>
          <ActivityIndicator size='large' color='#FF0266' />
        </View> : null}
        {mode === 'results' ?
              this.renderResults() : this.renderCameraCapture()}
      </View>
    );
  }
}

const styles = StyleSheet.create({
  sectionContainer: {
    marginTop: 32,
    paddingHorizontal: 24
  },
  centeredText: {
    textAlign: 'center',
    fontSize: 14,
  },
  sectionTitle: {
    fontSize: 24,
    fontWeight: '600',
    color: 'black',
    marginBottom: 6
  },
  loadingIndicator: {
    position: 'absolute',
    top: 20,
    right: 20,
    // flexDirection: 'row',
    // justifyContent: 'flex-end',
    zIndex: 200,
    // width: '100%'
  },
  cameraContainer: {
    display: 'flex',
    justifyContent: 'center',
    alignItems: 'center',
    width: '100%',
    height: '100%',
    backgroundColor: '#fff',
  },
  camera : {
    display: 'flex',
    width: '92%',
    height: '64%',
    backgroundColor: '#f0F',
    zIndex: 1,
    borderWidth: 20,
    borderRadius: 40,
    borderColor: '#f0f',
  },
  cameraControls: {
    display: 'flex',
    flexDirection: 'row',
    width: '92%',
    justifyContent: 'space-between',
    marginTop: 40,
    zIndex: 100,
    backgroundColor: 'transparent',
  },
  flipCameraBtn: {
    backgroundColor: '#424242',
    width: 75,
    height: 75,
    borderRadius:16,
    display: 'flex',
    flexDirection: 'row',
    justifyContent: 'center',
    alignItems: 'center',
  },
  takeImageBtn: {
    backgroundColor: '#FF0266',
    width: 75,
    height: 75,
    borderRadius:50,
    display: 'flex',
    flexDirection: 'row',
    justifyContent: 'center',
    alignItems: 'center',
  },
  cancelBtn: {
    backgroundColor: '#424242',
    width: 75,
    height: 75,
    borderRadius:4,
    display: 'flex',
    flexDirection: 'row',
    justifyContent: 'center',
    alignItems: 'center',
  },
  resultImageContainer : {
    width: '100%',
    height: '100%',
    padding:5,
    margin:0,
    backgroundColor: '#fff',
    zIndex: 1,
  },
  resultImage: {
    width: '98%',
    height: '98%',
  },
  styleImageContainer: {
    position:'absolute',
    width: 80,
    height: 150,
    bottom: 30,
    left: 20,
    zIndex: 10,
    borderRadius:10,
    backgroundColor: 'rgba(176, 222, 255, 0.5)',
    borderWidth: 1,
    borderColor: 'rgba(176, 222, 255, 0.7)',
  },
  contentImageContainer: {
    position:'absolute',
    width: 80,
    height: 150,
    bottom:30,
    right: 20,
    zIndex: 10,
    borderRadius:10,
    backgroundColor: 'rgba(255, 197, 161, 0.5)',
    borderWidth: 1,
    borderColor: 'rgba(255, 197, 161, 0.7)',
  },
  imagePreview: {
    width: 78,
    height: 148,
    borderRadius:10,
  },
  instructionText: {
    fontSize: 28,
    fontWeight:'bold',
    paddingLeft: 5
  },
  introText: {
    fontSize: 52,
    fontWeight:'bold',
    padding: 20,
    textAlign: 'left',
  }

});

style_transfer.ts;

import * as tf from '@tensorflow/tfjs';
import { fetch, bundleResourceIO } from '@tensorflow/tfjs-react-native'

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';

    // const STYLENET_URL = "http://instareportpro.com/saved_model_style_js/model.json";
    // const TRANSFORMNET_URL = "http://instareportpro.com/saved_model_transformer_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');

      // const modelJson = require('../ts_models/saved_model_style_js/model.json')
      // const modelWeight1 = require('../ts_models/saved_model_style_js/group1-shard1of3.bin')
      // const modelWeight2 = require('../ts_models/saved_model_style_js/group1-shard2of3.bin')
      // const modelWeight3 = require('../ts_models/saved_model_style_js/group1-shard3of3.bin')
      // this.styleNet = await tf.loadGraphModel(bundleResourceIO(modelJson, [
      //   modelWeight1, modelWeight2, modelWeight3]));
      //   console.log('stylenet loaded');

    }
  }

  async loadTransformerModel() {
    if (this.transformNet == null) {
      this.transformNet = await tf.loadGraphModel(TRANSFORMNET_URL);

      // const modelJson = require('../ts_models/saved_model_transformer_separable_js/model.json')
      // const modelWeight1 = require('../ts_models/saved_model_transformer_separable_js/group1-shard1of1.bin')
      // this.transformNet = await tf.loadGraphModel(bundleResourceIO(modelJson, [
      //   modelWeight1]));

        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');
      }

      console.log("produceStylized 111")
      const input = contentImage.toFloat().div(tf.scalar(255)).expandDims();
      console.log("produceStylized 222")
      const image: tf.Tensor4D =
          this.transformNet.predict([input, bottleneck]) as tf.Tensor4D;
          console.log("produceStylized 333")
      return image.mul(255).squeeze();
    });
  }

  public stylize(styleImage: tf.Tensor3D, contentImage: tf.Tensor3D):
      tf.Tensor3D {
          console.log("stylize 111")
    const start = Date.now();
    console.log("stylize 2222")
    // console.log(styleImage.shape, contentImage.shape);
    const styleRepresentation = this.predictStyleParameters(styleImage);
    console.log("stylize 3333")
    const stylized = this.produceStylized(contentImage, styleRepresentation);
    console.log("stylize 4444")
    tf.dispose([styleRepresentation]);
    console.log("stylize 5555")
    const end = Date.now();
    console.log("stylize 6666")
    console.log('stylization scheduled', end - start);
    return stylized;
  }
}

[tafsiri]

In this case you can probably use less memory by loading smaller images. Or you can resize large images to smaller ones using resizeBilinear. You may also need to call tf.dispose() on tensors once you are done using them to free up memory for subsequent operations.

[google-ml-butler[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. Thank you.

[google-ml-butler[bot]]

Closing as stale. Please @mention us if this needs more attention.

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No