NV21 Utils

NV21Utils.cpp

void convertNV21ToArray2d(JNIEnv* env, dlib::array2d<dlib::rgb_pixel>& out,
                          jbyteArray data, jint width, jint height) {
    jbyte* yuv = env->GetByteArrayElements(data, 0);
    int frameSize = width * height;
    int y, u, v, uvIndex;
    int r, g, b;

    out.set_size((long) height, (long) width);

    for(int row = 0; row < height; row++) {
        for(int column = 0; column < width; column++) {
            uvIndex = frameSize + (row >> 1) * width + (column & ~1);
            y = 0xff & yuv[row * width + column] - 16;
            v = 0xff & yuv[uvIndex] - 128;
            u = 0xff & yuv[uvIndex+1] - 128;

            y = y < 0 ? 0 : 1164 * y;

            r = y + 1596 * v;
            g = y - 813 * v - 391 * u;
            b = y + 2018 * u;

            out[row][column].red = (unsigned char)(r < 0 ? 0 : r > 255000 ? 255 : r/1000);
            out[row][column].green = (unsigned char)(g < 0 ? 0 : g > 255000 ? 255 : g/1000);
            out[row][column].blue = (unsigned char)(b < 0 ? 0 : b > 255000 ? 255 : b/1000);
        }
    }
}

extern "C" JNIEXPORT jbyteArray JNICALL
JNI_METHOD(downScaleNV21)(JNIEnv* env, jclass obj,
                          jbyteArray data, jint width, jint height) {

    int width_ds = width >> 1;
    int height_ds = height >> 1;

    int frameSize = width * height;
    int frameSize_ds = width_ds * height_ds;

    int yuv_ds_size = (frameSize_ds * 3) >> 1;

    jbyte* yuv = env->GetByteArrayElements(data, 0);

    jbyte* yuv_ds = new jbyte[yuv_ds_size];
    jbyteArray out = env->NewByteArray(yuv_ds_size);

    int y1, y2, y3, y4;
    int u, v;
    int uvIndex, uvIndex1, uvIndex2, uvIndex3, uvIndex4;
    int row_2, column_2;

    for(int row_0 = 0, row_ds = 0; row_0 < height; row_0 += 4, row_ds += 2) {
        row_2 = row_0 + 2;
        for(int column_0 = 0, column_ds = 0; column_0 < width; column_0 += 4, column_ds += 2) {
            //      0     1     2     3
            //    _______________________
            // 0 | y1a | y1b | y2a | y2b |
            //   |-----+-----|-----+-----|          _____ _____
            // 1 | y1c | y1d | y2c | y2d |         | y1a | y2a |
            //    -----------+-----------     ->    -----+-----
            // 2 | y3a | y3b | y4a | y4b |         | y3a | y4a |
            //   |-----+-----|-----+-----|          ----- -----
            // 3 | y3c | y3d | y4c | y4d |
            //    ----------- -----------

            column_2 = column_0 + 2;

            y1 = yuv[row_0 * width + column_0];

            y2 = yuv[row_0 * width + column_2];

            y3 = yuv[row_2 * width + column_0];

            y4 = yuv[row_2 * width + column_2];

            uvIndex1 = frameSize +
                       (row_0 >> 1) * width + (column_0 & ~1);         // v_index of y1a
            uvIndex = frameSize_ds +
                      (row_ds >> 1) * width_ds + (column_ds & ~1);    // v_index of y1

            v = yuv[uvIndex1];
            u = yuv[uvIndex1+1];

            yuv_ds[row_ds * width_ds + column_ds] = (jbyte) y1;
            yuv_ds[row_ds * width_ds + (column_ds + 1)] = (jbyte) y2;
            yuv_ds[(row_ds + 1) * width_ds + column_ds] = (jbyte) y3;
            yuv_ds[(row_ds + 1) * width_ds + (column_ds + 1)] = (jbyte) y4;
            yuv_ds[uvIndex] = (jbyte) v;
            yuv_ds[uvIndex + 1] = (jbyte) u;
        }
    }
    env->SetByteArrayRegion(out, 0, yuv_ds_size, yuv_ds);
    delete[] yuv_ds;
    return out;
}

extern "C" JNIEXPORT jbyteArray JNICALL
JNI_METHOD(rotateNV21)(JNIEnv* env, jclass obj,
                       jbyteArray data, jint width, jint height, jint rotation) {

    int frameSize = width * height;
    int yuv_size = (frameSize * 3) >> 1;
    jbyte* yuv = env->GetByteArrayElements(data, 0);
    jbyte* yuv_rotated = new jbyte[yuv_size];

    jbyteArray out = env->NewByteArray(yuv_size);

    int row_r, column_r;
    int width_r;
    int height_r;

    if(rotation == 180) {
        width_r = width;
        height_r = height;
    } else {
        width_r = height;
        height_r = width;
    }

    int uvIndex_r, uvIndex;
    for(int row = 0; row < height; row++) {
        for (int column = 0; column < width; column++) {
            if(rotation == 90) {
                column_r = width_r - row - 1;
                row_r = column;
            } else if(rotation == 180) {
                column_r = width_r - column - 1;
                row_r = height_r - row - 1;
            } else { // assume its 270
                column_r = row;
                row_r = height_r - column - 1;
            }

            yuv_rotated[row_r * width_r + column_r] = yuv[row * width + column];

            if((row & 1) == 0 &&
               (column & 1) == 0) {
                uvIndex = frameSize + (row >> 1) * width + (column & ~1);
                uvIndex_r = frameSize + (row_r >> 1) * width_r + (column_r & ~1);

                yuv_rotated[uvIndex_r] = yuv[uvIndex];
                yuv_rotated[uvIndex_r + 1] = yuv[uvIndex + 1];
            }
        }
    }
    env->SetByteArrayRegion(out, 0, yuv_size, yuv_rotated);
    delete[] yuv_rotated;
    return out;
}

extern "C" JNIEXPORT jbyteArray JNICALL
JNI_METHOD(flipNV21)(JNIEnv* env, jclass obj,
                     jbyteArray data, jint width, jint height, jint direction) {
    int frameSize = width * height;
    int yuv_size = (frameSize * 3) >> 1;
    jbyte* yuv = env->GetByteArrayElements(data, 0);

    jbyteArray out = env->NewByteArray(yuv_size);

    int column_f;
    int row_f;
    int index, index_f;
    int uvIndex, uvIndex_f;
    int height_mid;
    int width_mid;
    jbyte tmp;

    if(direction == 0) {    // Horizontal
        width_mid = ((width & 1) == 0) ? (width / 2) : (width / 2 + 1);
        height_mid = height;
    } else {                // Vertical
        height_mid = ((height & 1) == 0) ? (height / 2) : (height / 2 + 1);
        width_mid = width;
    }

    for(int row = 0; row < height_mid; row++) {
        for (int column = 0; column < width_mid; column++) {
            if(direction == 0) {    // Horizontal
                column_f = width - column - 1;
                row_f = row;
            } else {                // Vertical
                column_f = column;
                row_f = height - row - 1;
            }

            index = row * width + column;
            index_f = row_f * width + column_f;

            tmp = yuv[index];
            yuv[index] = yuv[index_f];
            yuv[index_f] = tmp;

            if((row & 1) == 0 &&
               (column & 1) == 0) {
                uvIndex = frameSize + (row >> 1) * width + (column & ~1);
                uvIndex_f = frameSize + (row_f >> 1) * width + (column_f & ~1);

                tmp = yuv[uvIndex];
                yuv[uvIndex] = yuv[uvIndex_f];
                yuv[uvIndex_f] = tmp;

                tmp = yuv[uvIndex + 1];
                yuv[uvIndex + 1] = yuv[uvIndex_f + 1];
                yuv[uvIndex_f + 1] = tmp;
            }
        }
    }
    env->SetByteArrayRegion(out, 0, yuv_size, yuv);
    return out;
}