shaqian / TF-Unity

54 stars 11 forks source link
tensorflow unity

TF-Unity

Scripts for accessing TensorFlow API in Unity3d. Supports image classification and object detection (SSD or Tiny YOLOv2).

Installation

1) Install Unity TensorFlow Plugin.

For details refer to Using TensorFlowSharp in Unity (Experimental).

  1. Download and import Unity TensorFlow Plugin (Download here).

  2. Go to Edit -> Player Settings and add ENABLE_TENSORFLOW to the Scripting Define Symbols for each type of device you want to use (PC, Mac and Linux Standalone, iOS or Android).

  3. iOS additional instructions for building

2) Copy the TensorFlow folder in this repo to the Assets folder of the Unity project.

3) Copy the Models folder in this repo to the Assets folder of the Unity project.

Usage

Image Classification

var classifier = new Classifier(
  model,                                       // TextAsset
  labels,                                      // TextAsset
  input: "input",                              // Optional, defaults to "input"
  output: "MobilenetV1/Predictions/Reshape_1", // Optional, defaults to "output"
  height: 300,                                 // Optional, defaults to 224
  width: 300,                                  // Optional, defaults to 224
  mean: 127.5f,                                // Optional, defaults to 127.5f
  std: 127.5f                                  // Optional, defaults to 127.5f
);
var outputs = classifier.Classify(
  texture,             // Texture2D
  numResults: 1,       // Optional, defauls to 5
  angle: 90,           // Optional, defauls to 0
  threshold: 0.05f,    // Optional, defauls to 0.1f
  flip: Flip.VERTICAL, // Optional, default to Flip.NONE
);
[ // List
  { // KeyValuePair
    "apple": 0.1 // string: float
  },
  ......
]

Object Detection

var detector = new Detector(TextAsset 
  modelFile,                   // TextAsset
  labelFile,                   // TextAsset
  model: DetectionModels.YOLO, // Optional, defaults to DetectionModels.SSD
  input: "image_tensor",       // Optional, defaults to "input"
  output: "output",            // Optional, defaults to "output"
  height: 300,                 // Optional, defaults to 300
  width: 300,                  // Optional, defaults to 300
  mean: 127.5f,                // Optional, defaults to 127.5f
  std: 127.5f                  // Optional, defaults to 127.5f
  blockSize: 32,               // Used in YOLO only, defaults to 32
  numBoxesPerBlock: 5,         // Used in YOLO only, defaults to 5
  anchors: []                  // Used in YOLO only, defaults to [0.57273,0.677385,1.87446,2.06253,3.33843,5.47434,7.88282,3.52778,9.77052,9.16828]
);
var outputs = detector.detect(
  texture,               // Texture2D
  numResultsPerClass: 5, // Optional, defauls to 1
  angle: 90,             // Optional, defauls to 0
  threshold: 0.5f,       // Optional, defauls to 0.2f
  flip: Flip.VERTICAL,   // Optional, default to Flip.NONE
);

x, y, w, h are between [0, 1]. You can scale x, w by the width and y, h by the height of the image.

[ // List
  { // Dictionary
    "detectedClass": "hot dog", // string: string
    "confidenceInClass": 0.123, // string: float
    "rect": { // Dictionary
      "x": 0.15, // string: float
      "y": 0.33,
      "w": 0.80,
      "h": 0.27
    }
  },
  ......
]

Utils

Utils.DrawOutput(
  outputs, 
  Screen.width,  // Scale width
  Screen.height, // Scale height
  Color.yellow
);

Example

Refer to TF-Unity-ARFoundation.