-
Notifications
You must be signed in to change notification settings - Fork 6
/
YuvToRgbConverter.kt
163 lines (143 loc) · 6.16 KB
/
YuvToRgbConverter.kt
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
152
153
154
155
156
157
158
159
160
161
162
163
package com.appinmotion.gpuimage
import android.content.Context
import android.graphics.Bitmap
import android.graphics.ImageFormat
import android.graphics.Rect
import android.media.Image
import android.renderscript.Allocation
import android.renderscript.Element
import android.renderscript.RenderScript
import android.renderscript.ScriptIntrinsicYuvToRGB
import java.nio.ByteBuffer
/**
* Helper class used to efficiently convert a [Media.Image] object from
* YUV_420_888 format to an RGB [Bitmap] object.
*
* Copied from https://github.com/owahltinez/camerax-tflite/blob/master/app/src/main/java/com/android/example/camerax/tflite/YuvToRgbConverter.kt
*
* The [yuvToRgb] method is able to achieve the same FPS as the CameraX image
* analysis use case at the default analyzer resolution, which is 30 FPS with
* 640x480 on a Pixel 3 XL device.
*/
class YuvToRgbConverter(context: Context) {
private val rs = RenderScript.create(context)
private val scriptYuvToRgb = ScriptIntrinsicYuvToRGB.create(rs, Element.U8_4(rs))
private var pixelCount: Int = -1
private lateinit var yuvBuffer: ByteBuffer
private lateinit var inputAllocation: Allocation
private lateinit var outputAllocation: Allocation
@Synchronized
fun yuvToRgb(image: Image, output: Bitmap) {
// Ensure that the intermediate output byte buffer is allocated
if (!::yuvBuffer.isInitialized) {
pixelCount = image.cropRect.width() * image.cropRect.height()
yuvBuffer = ByteBuffer.allocateDirect(
pixelCount * ImageFormat.getBitsPerPixel(ImageFormat.YUV_420_888) / 8
)
}
// Get the YUV data in byte array form
imageToByteBuffer(image, yuvBuffer)
// Ensure that the RenderScript inputs and outputs are allocated
if (!::inputAllocation.isInitialized) {
inputAllocation = Allocation.createSized(rs, Element.U8(rs), yuvBuffer.array().size)
}
if (!::outputAllocation.isInitialized) {
outputAllocation = Allocation.createFromBitmap(rs, output)
}
// Convert YUV to RGB
inputAllocation.copyFrom(yuvBuffer.array())
scriptYuvToRgb.setInput(inputAllocation)
scriptYuvToRgb.forEach(outputAllocation)
outputAllocation.copyTo(output)
}
private fun imageToByteBuffer(image: Image, outputBuffer: ByteBuffer) {
assert(image.format == ImageFormat.YUV_420_888)
val imageCrop = image.cropRect
val imagePlanes = image.planes
val rowData = ByteArray(imagePlanes.first().rowStride)
imagePlanes.forEachIndexed { planeIndex, plane ->
// How many values are read in input for each output value written
// Only the Y plane has a value for every pixel, U and V have half the resolution i.e.
//
// Y Plane U Plane V Plane
// =============== ======= =======
// Y Y Y Y Y Y Y Y U U U U V V V V
// Y Y Y Y Y Y Y Y U U U U V V V V
// Y Y Y Y Y Y Y Y U U U U V V V V
// Y Y Y Y Y Y Y Y U U U U V V V V
// Y Y Y Y Y Y Y Y
// Y Y Y Y Y Y Y Y
// Y Y Y Y Y Y Y Y
val outputStride: Int
// The index in the output buffer the next value will be written at
// For Y it's zero, for U and V we start at the end of Y and interleave them i.e.
//
// First chunk Second chunk
// =============== ===============
// Y Y Y Y Y Y Y Y U V U V U V U V
// Y Y Y Y Y Y Y Y U V U V U V U V
// Y Y Y Y Y Y Y Y U V U V U V U V
// Y Y Y Y Y Y Y Y U V U V U V U V
// Y Y Y Y Y Y Y Y
// Y Y Y Y Y Y Y Y
// Y Y Y Y Y Y Y Y
var outputOffset: Int
when (planeIndex) {
0 -> {
outputStride = 1
outputOffset = 0
}
1 -> {
outputStride = 2
outputOffset = pixelCount + 1
}
2 -> {
outputStride = 2
outputOffset = pixelCount
}
else -> {
// Image contains more than 3 planes, something strange is going on
return@forEachIndexed
}
}
val buffer = plane.buffer
val rowStride = plane.rowStride
val pixelStride = plane.pixelStride
// We have to divide the width and height by two if it's not the Y plane
val planeCrop = if (planeIndex == 0) {
imageCrop
} else {
Rect(
imageCrop.left / 2,
imageCrop.top / 2,
imageCrop.right / 2,
imageCrop.bottom / 2
)
}
val planeWidth = planeCrop.width()
val planeHeight = planeCrop.height()
buffer.position(rowStride * planeCrop.top + pixelStride * planeCrop.left)
for (row in 0 until planeHeight) {
val length: Int
if (pixelStride == 1 && outputStride == 1) {
// When there is a single stride value for pixel and output, we can just copy
// the entire row in a single step
length = planeWidth
buffer.get(outputBuffer.array(), outputOffset, length)
outputOffset += length
} else {
// When either pixel or output have a stride > 1 we must copy pixel by pixel
length = (planeWidth - 1) * pixelStride + 1
buffer.get(rowData, 0, length)
for (col in 0 until planeWidth) {
outputBuffer.array()[outputOffset] = rowData[col * pixelStride]
outputOffset += outputStride
}
}
if (row < planeHeight - 1) {
buffer.position(buffer.position() + rowStride - length)
}
}
}
}
}