/
SegmentImageTask.java
151 lines (119 loc) · 5.75 KB
/
SegmentImageTask.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
package com.qdn.segmentation.tasks;
import android.content.Context;
import android.graphics.Bitmap;
import android.os.AsyncTask;
import com.qdn.segmentation.Utils.Logger;
import com.qualcomm.qti.snpe.FloatTensor;
import com.qualcomm.qti.snpe.NeuralNetwork;
import com.qdn.segmentation.Helpers.BitmapToFloatArrayHelper;
import com.qdn.segmentation.Utils.Constants;
import com.qdn.segmentation.Interfaces.IBitmapLoader;
import java.util.HashMap;
import java.util.Map;
import static com.qdn.segmentation.Utils.Constants.MNETSSD_INPUT_LAYER;
import static com.qdn.segmentation.Utils.Constants.MNETSSD_OUTPUT_LAYER;
public class SegmentImageTask extends AsyncTask<Void, Void, Bitmap> {
private static String TAG = SegmentImageTask.class.getSimpleName();
private static int MNETSSD_NUM_BOXES = Constants.BITMAP_WIDTH * Constants.BITMAP_WIDTH;
private final float[] floatOutput = new float[MNETSSD_NUM_BOXES];
private Map<String, FloatTensor> mOutputs;
private BitmapToFloatArrayHelper mBitmapToFloatHelper;
public int originalBitmapW, originalBitmapH;
private int[] mInputTensorShapeHWC;
private FloatTensor mInputTensorReused;
private Map<String, FloatTensor> mInputTensorsMap;
private NeuralNetwork mNeuralnetwork;
private Bitmap mScaledBitmap, mOutputBitmap;
private IBitmapLoader mCallbackBitmapLoader;
public int getInputTensorWidth() {
return mInputTensorShapeHWC == null ? 0 : mInputTensorShapeHWC[1];
}
public int getInputTensorHeight() {
return mInputTensorShapeHWC == null ? 0 : mInputTensorShapeHWC[2];
}
public SegmentImageTask(final Context context,
final NeuralNetwork neuralNetwork,
final Bitmap bitmap,
final IBitmapLoader iBitmapLoader) {
this.mNeuralnetwork = neuralNetwork;
this.mCallbackBitmapLoader = iBitmapLoader;
mBitmapToFloatHelper = new BitmapToFloatArrayHelper();
originalBitmapH = bitmap.getHeight();
originalBitmapW = bitmap.getWidth();
mScaledBitmap = Bitmap.createScaledBitmap(bitmap, Constants.BITMAP_WIDTH, Constants.BITMAP_HEIGHT, false);
}
@Override
protected Bitmap doInBackground(Void... params) {
return deeplabV3Inference();
}
@Override
protected void onPostExecute(Bitmap bitmap) {
super.onPostExecute(bitmap);
Logger.d(TAG, "onPostExecute " + bitmap);
mCallbackBitmapLoader.loadResultBitmap(bitmap);
}
public Bitmap deeplabV3Inference() {
try {
mInputTensorShapeHWC = mNeuralnetwork.getInputTensorsShapes().get(MNETSSD_INPUT_LAYER);
// allocate the single input tensor
mInputTensorReused = mNeuralnetwork.createFloatTensor(mInputTensorShapeHWC);
// add it to the map of inputs, even if it's a single input
mInputTensorsMap = new HashMap<>();
mInputTensorsMap.put(MNETSSD_INPUT_LAYER, mInputTensorReused);
// execute the inference, and get 3 tensors as outputs
mOutputs = inferenceOnBitmap(mScaledBitmap);
if (mOutputs != null) {
MNETSSD_NUM_BOXES = mOutputs.get(MNETSSD_OUTPUT_LAYER).getSize();
// convert tensors to boxes - Note: Optimized to read-all upfront
mOutputs.get(MNETSSD_OUTPUT_LAYER).read(floatOutput, 0, MNETSSD_NUM_BOXES);
//for black/white image
int w = mScaledBitmap.getWidth();
int h = mScaledBitmap.getHeight();
int b = 0xFF;
int out = 0xFF;
for (int y = 0; y < h; y++) {
for (int x = 0; x < w; x++) {
b = b & mScaledBitmap.getPixel(x, y);
for (int i = 1; i <= 3 && floatOutput[y * w + x] != 15; i++) {
out = out << (8) | b;
}
mScaledBitmap.setPixel(x, y, floatOutput[y * w + x] != 15 ? out : mScaledBitmap.getPixel(x, y));
out = 0xFF;
b = 0xFF;
}
}
mOutputBitmap = Bitmap.createScaledBitmap(mScaledBitmap, originalBitmapW,
originalBitmapH, true);
Logger.d(TAG, mOutputBitmap.getWidth() + "");
}
} catch (Exception e) {
e.printStackTrace();
}
return mOutputBitmap;
}
/* Generic functions, for typical image models */
private Map<String, FloatTensor> inferenceOnBitmap(Bitmap scaledBitmap) {
final Map<String, FloatTensor> outputs;
try {
// safety check
if (mNeuralnetwork == null || mInputTensorReused == null || scaledBitmap.getWidth() != getInputTensorWidth() || scaledBitmap.getHeight() != getInputTensorHeight()) {
Logger.d("SNPEHelper", "No NN loaded, or image size different than tensor size");
return null;
}
// [0.3ms] Bitmap to RGBA byte array (size: 300*300*3 (RGBA..))
mBitmapToFloatHelper.bitmapToBuffer(scaledBitmap);
// [2ms] Pre-processing: Bitmap (300,300,4 ints) -> Float Input Tensor (300,300,3 floats)
final float[] inputFloatsHW3 = mBitmapToFloatHelper.bufferToNormalFloatsBGR();
if (mBitmapToFloatHelper.isFloatBufferBlack())
return null;
mInputTensorReused.write(inputFloatsHW3, 0, inputFloatsHW3.length, 0, 0);
// [31ms on GPU16, 50ms on GPU] execute the inference
outputs = mNeuralnetwork.execute(mInputTensorsMap);
} catch (Exception e) {
e.printStackTrace();
Logger.d("SNPEHelper", e.getCause() + "");
return null;
}
return outputs;
}
}