Skip to content

Commit 65778b9

Browse files
jeffxtangbrianjo
andauthored
iOS and Android image segmentation tutorials; minor recipes code and typo fix for 1.7 (#1213)
* pt mobile script and optimize recipe * 1 pt mobile new recipes summary and 5 recipes * updated recipes_index.rst * thumbnail png fix for ios recipe in recipes_index.rst * edits based on feedback * recipes code and typo fix on 1.7 * merge from master * fixed merge errors * added limitation warning for dynamic quantization * link fix for quantization and model prep for ios recipes * deeplabv3_on_ios tutorial * more edits on the ios segmentation tutorial * image segmentation android tutorial * ios segmentation tutorial update * code refactoring on code coloring; added reviewed by * updated android deeplabv3 tutorial * further update on the iOS and Android segmentation tutorials * minor and maybe final edits for two tutorials Co-authored-by: Brian Johnson <brianjo@fb.com>
1 parent 74b7663 commit 65778b9

10 files changed

+496
-19
lines changed

_static/img/deeplabv3_android.png

371 KB
Loading

_static/img/deeplabv3_android2.png

54.7 KB
Loading

_static/img/deeplabv3_ios.png

381 KB
Loading

_static/img/deeplabv3_ios2.png

42.6 KB
Loading
Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
Image Segmentation DeepLabV3 on Android
2+
=================================================
3+
4+
**Author**: `Jeff Tang <https://github.com/jeffxtang>`_
5+
6+
**Reviewed by**: `Jeremiah Chung <https://github.com/jeremiahschung>`_
7+
8+
Introduction
9+
------------
10+
11+
Semantic image segmentation is a computer vision task that uses semantic labels to mark specific regions of an input image. The PyTorch semantic image segmentation `DeepLabV3 model <https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101>`_ can be used to label image regions with `20 semantic classes <http://host.robots.ox.ac.uk:8080/pascal/VOC/voc2007/segexamples/index.html>`_ including, for example, bicycle, bus, car, dog, and person. Image segmentation models can be very useful in applications such as autonomous driving and scene understanding.
12+
13+
In this tutorial, we will provide a step-by-step guide on how to prepare and run the PyTorch DeepLabV3 model on Android, taking you from the beginning of having a model you may want to use on Android to the end of having a complete Android app using the model. We will also cover practical and general tips on how to check if your next favorable pre-trained PyTorch models can run on Android, and how to avoid pitfalls.
14+
15+
.. note:: Before going through this tutorial, you should check out `PyTorch Mobile for Android <https://pytorch.org/mobile/android/>`_ and give the PyTorch Android `HelloWorld <https://github.com/pytorch/android-demo-app/tree/master/HelloWorldApp>`_ example app a quick try. This tutorial will go beyond the image classification model, usually the first kind of model deployed on mobile. The complete code repo for this tutorial is available `here <https://github.com/pytorch/android-demo-app/ImageSegmentation>`_.
16+
17+
Learning Objectives
18+
-------------------
19+
20+
In this tutorial, you will learn how to:
21+
22+
1. Convert the DeepLabV3 model for Android deployment.
23+
24+
2. Get the output of the model for the example input image in Python and compare it to the output from the Android app.
25+
26+
3. Build a new Android app or reuse an Android example app to load the converted model.
27+
28+
4. Prepare the input into the format that the model expects and process the model output.
29+
30+
5. Complete the UI, refactor, build and run the app to see image segmentation in action.
31+
32+
Pre-requisites
33+
---------------
34+
35+
* PyTorch 1.6 or 1.7
36+
37+
* torchvision 0.7 or 0.8
38+
39+
* Android Studio 3.5.1 or above with NDK installed
40+
41+
Steps
42+
---------
43+
44+
1. Convert the DeepLabV3 model for Android deployment
45+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
46+
47+
The first step to deploying a model on Android is to convert the model into the `TorchScript <https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html>`_ format.
48+
49+
.. note::
50+
Not all PyTorch models can be converted to TorchScript at this time because a model definition may use language features that are not in TorchScript, which is a subset of Python. See the `Script and Optimize Recipe <../recipes/script_optimized.html>`_ for more details.
51+
52+
Simply run the script below to generate the scripted model `deeplabv3_scripted.pt`:
53+
54+
::
55+
56+
import torch
57+
58+
# use deeplabv3_resnet50 instead of resnet101 to reduce the model size
59+
model = torch.hub.load('pytorch/vision:v0.7.0', 'deeplabv3_resnet50', pretrained=True)
60+
model.eval()
61+
62+
scriptedm = torch.jit.script(model)
63+
torch.jit.save(scriptedm, "deeplabv3_scripted.pt")
64+
65+
The size of the generated `deeplabv3_scripted.pt` model file should be around 168MB. Ideally, a model should also be quantized for significant size reduction and faster inference before being deployed on an Android app. To have a general understanding of quantization, see the `Quantization Recipe <../recipes/quantization.html>`_ and the resource links there. We will cover in detail how to correctly apply a quantization workflow called Post Training `Static Quantization <https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html>`_ to the DeepLabV3 model in a future tutorial or recipe.
66+
67+
2. Get example input and output of the model in Python
68+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
69+
70+
Now that we have a scripted PyTorch model, let's test with some example inputs to make sure the model works correctly on Android. First, let's write a Python script that uses the model to make inferences and examine inputs and outputs. For this example of the DeepLabV3 model, we can reuse the code in Step 1 and in the `DeepLabV3 model hub site <https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101>`_. Add the following code snippet to the code above:
71+
72+
::
73+
74+
from PIL import Image
75+
from torchvision import transforms
76+
input_image = Image.open("deeplab.jpg")
77+
preprocess = transforms.Compose([
78+
transforms.ToTensor(),
79+
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
80+
])
81+
82+
input_tensor = preprocess(input_image)
83+
input_batch = input_tensor.unsqueeze(0)
84+
with torch.no_grad():
85+
output = model(input_batch)['out'][0]
86+
87+
print(input_batch.shape)
88+
print(output.shape)
89+
90+
Download `deeplab.jpg` from `here <https://github.com/jeffxtang/android-demo-app/blob/new_demo_apps/ImageSegmentation/app/src/main/assets/deeplab.jpg>`_, then run the script above and you will see the shapes of the input and output of the model:
91+
92+
::
93+
94+
torch.Size([1, 3, 400, 400])
95+
torch.Size([21, 400, 400])
96+
97+
So if you provide the same image input `deeplab.jpg` of size 400x400 to the model on Android, the output of the model should have the size [21, 400, 400]. You should also print out at least the beginning parts of the actual data of the input and output, to be used in Step 4 below to compare with the actual input and output of the model when running in the Android app.
98+
99+
3. Build a new Android app or reuse an example app and load the model
100+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
101+
102+
First, follow Step 3 of the `Model Preparation for Android recipe <../recipes/model_preparation_android.html#add-the-model-and-pytorch-library-on-android>`_ to use our model in an Android Studio project with PyTorch Mobile enabled. Because both DeepLabV3 used in this tutorial and MobileNet v2 used in the PyTorch HelloWorld Android example are computer vision models, you can also get the `HelloWorld example repo <https://github.com/pytorch/android-demo-app/tree/master/HelloWorldApp>`_ to make it easier to modify the code that loads the model and processes the input and output. The main goal in this step and Step 4 is to make sure the model `deeplabv3_scripted.pt` generated in Step 1 can indeed work correctly on Android.
103+
104+
Now let's add `deeplabv3_scripted.pt` and `deeplab.jpg` used in Step 2 to the Android Studio project and modify the `onCreate` method in the `MainActivity` to resemble:
105+
106+
.. code-block:: java
107+
108+
Module module = null;
109+
try {
110+
module = Module.load(assetFilePath(this, "deeplabv3_scripted.pt"));
111+
} catch (IOException e) {
112+
Log.e("ImageSegmentation", "Error loading model!", e);
113+
finish();
114+
}
115+
116+
Then set a breakpoint at the line `finish()` and build and run the app. If the app doesn't stop at the breakpoint, it means that the scripted model in Step 1 has been successfully loaded on Android.
117+
118+
4. Process the model input and output for model inference
119+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
120+
121+
After the model loads in the previous step, let's verify that it works with expected inputs and can generate expected outputs. As the model input for the DeepLabV3 model is an image the same as that of the MobileNet v2 in the HelloWorld example, we will reuse some of the code in the `MainActivity.java <https://github.com/pytorch/android-demo-app/blob/master/HelloWorldApp/app/src/main/java/org/pytorch/helloworld/MainActivity.java>`_ file from HelloWorld for input processing. Replace the code snippet between `line 50 <https://github.com/pytorch/android-demo-app/blob/master/HelloWorldApp/app/src/main/java/org/pytorch/helloworld/MainActivity.java#L50>`_ and 73 in `MainActivity.java` with the following code:
122+
123+
.. code-block:: java
124+
125+
final Tensor inputTensor = TensorImageUtils.bitmapToFloat32Tensor(bitmap,
126+
TensorImageUtils.TORCHVISION_NORM_MEAN_RGB,
127+
TensorImageUtils.TORCHVISION_NORM_STD_RGB);
128+
final float[] inputs = inputTensor.getDataAsFloatArray();
129+
130+
Map<String, IValue> outTensors =
131+
module.forward(IValue.from(inputTensor)).toDictStringKey();
132+
133+
// the key "out" of the output tensor contains the semantic masks
134+
// see https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101
135+
final Tensor outputTensor = outTensors.get("out").toTensor();
136+
final float[] outputs = outputTensor.getDataAsFloatArray();
137+
138+
int width = bitmap.getWidth();
139+
int height = bitmap.getHeight();
140+
141+
.. note::
142+
The model output is a dictionary for the DeepLabV3 model so we use `toDictStringKey` to correctly extract the result. For other models, the model output may also be a single tensor or a tuple of tensors, among other things.
143+
144+
With the code changes shown above, you can set breakpoints after `final float[] inputs` and `final float[] outputs`, which populate the input tensor and output tensor data to float arrays for easy debugging. Run the app and when it stops at the breakpoints, compare the numbers in `inputs` and `outputs` with the model input and output data you see in Step 2 to see if they match. For the same inputs to the models running on Android and Python, you should get the same outputs.
145+
146+
.. warning::
147+
You may see different model outputs with the same image input when running on an Android emulator due to some Android emulator's floating point implementation issue. So it is best to test the app on a real Android device.
148+
149+
All we have done so far is to confirm that the model of our interest can be scripted and run correctly in our Android app as in Python. The steps we walked through so far for using a model in an iOS app consumes the bulk, if not most, of our app development time, similar to how data preprocessing is the heaviest lift for a typical machine learning project.
150+
151+
5. Complete the UI, refactor, build and run the app
152+
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
153+
154+
Now we are ready to complete the app and the UI to actually see the processed result as a new image. The output processing code should be like this, added to the end of the code snippet in Step 4:
155+
156+
.. code-block:: java
157+
158+
int[] intValues = new int[width * height];
159+
// go through each element in the output of size [WIDTH, HEIGHT] and
160+
// set different color for different classnum
161+
for (int j = 0; j < width; j++) {
162+
for (int k = 0; k < height; k++) {
163+
// maxi: the index of the 21 CLASSNUM with the max probability
164+
int maxi = 0, maxj = 0, maxk = 0;
165+
double maxnum = -100000.0;
166+
for (int i=0; i < CLASSNUM; i++) {
167+
if (outputs[i*(width*height) + j*width + k] > maxnum) {
168+
maxnum = outputs[i*(width*height) + j*width + k];
169+
maxi = i; maxj = j; maxk= k;
170+
}
171+
}
172+
// color coding for person (red), dog (green), sheep (blue)
173+
// black color for background and other classes
174+
if (maxi == PERSON)
175+
intValues[maxj*width + maxk] = 0xFFFF0000; // red
176+
else if (maxi == DOG)
177+
intValues[maxj*width + maxk] = 0xFF00FF00; // green
178+
else if (maxi == SHEEP)
179+
intValues[maxj*width + maxk] = 0xFF0000FF; // blue
180+
else
181+
intValues[maxj*width + maxk] = 0xFF000000; // black
182+
}
183+
}
184+
185+
The constants used in the code above are defined in the beginning of the class `MainActivity`:
186+
187+
.. code-block:: java
188+
189+
private static final int CLASSNUM = 21;
190+
private static final int DOG = 12;
191+
private static final int PERSON = 15;
192+
private static final int SHEEP = 17;
193+
194+
195+
The implementation here is based on the understanding of the DeepLabV3 model which outputs a tensor of size [21, width, height] for an input image of width*height. Each element in the width*height output array is a value between 0 and 20 (for a total of 21 semantic labels described in Introduction) and the value is used to set a specific color. Color coding of the segmentation here is based on the class with the highest probability, and you can extend the color coding for all classes in your own dataset.
196+
197+
After the output processing, you will also need to call the code below to render the RGB `intValues` array to a bitmap instance `outputBitmap` before displaying it on an `ImageView`:
198+
199+
.. code-block:: java
200+
201+
Bitmap bmpSegmentation = Bitmap.createScaledBitmap(bitmap, width, height, true);
202+
Bitmap outputBitmap = bmpSegmentation.copy(bmpSegmentation.getConfig(), true);
203+
outputBitmap.setPixels(intValues, 0, outputBitmap.getWidth(), 0, 0,
204+
outputBitmap.getWidth(), outputBitmap.getHeight());
205+
imageView.setImageBitmap(outputBitmap);
206+
207+
The UI for this app is also similar to that for HelloWorld, except that you do not need the `TextView` to show the image classification result. You can also add two buttons `Segment` and `Restart` as shown in the code repo to run the model inference and to show back the original image after the segmentation result is shown.
208+
209+
Now when you run the app on an Android emulator or preferably an actual device, you will see screens like the following:
210+
211+
.. image:: /_static/img/deeplabv3_android.png
212+
:width: 300 px
213+
.. image:: /_static/img/deeplabv3_android2.png
214+
:width: 300 px
215+
216+
217+
Recap
218+
--------
219+
220+
In this tutorial, we described what it takes to convert a pre-trained PyTorch DeepLabV3 model for Android and how to make sure the model can run successfully on Android. Our focus was to help you understand the process of confirming that a model can indeed run on Android. The complete code repo is available `here <https://github.com/pytorch/android-demo-app/ImageSegmentation>`_.
221+
222+
More advanced topics such as quantization and using models via transfer learning or of your own on Android will be covered soon in future demo apps and tutorials.
223+
224+
225+
Learn More
226+
------------
227+
228+
1. `PyTorch Mobile site <https://pytorch.org/mobile>`_
229+
2. `DeepLabV3 model <https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101>`_
230+
3. `DeepLabV3 paper <https://arxiv.org/pdf/1706.05587.pdf>`_

0 commit comments

Comments
 (0)