stagefright: use colr box for color aspects

Bug: 25684127
Change-Id: I105294ce0d41ac58d80a2a4a74d35a9b78536790
diff --git a/include/media/stagefright/foundation/ALookup.h b/include/media/stagefright/foundation/ALookup.h
index d8af407..5a68806 100644
--- a/include/media/stagefright/foundation/ALookup.h
+++ b/include/media/stagefright/foundation/ALookup.h
@@ -27,14 +27,14 @@
 struct ALookup {
     ALookup(std::initializer_list<std::pair<T, U>> list);
 
-    bool lookup(const T& from, U *to);
-    bool rlookup(const U& from, T *to);
+    bool lookup(const T& from, U *to) const;
+    bool rlookup(const U& from, T *to) const;
 
     template<typename V, typename = typename std::enable_if<!std::is_same<T, V>::value>::type>
-    inline bool map(const T& from, V *to) { return lookup(from, to); }
+    inline bool map(const T& from, V *to) const { return lookup(from, to); }
 
     template<typename V, typename = typename std::enable_if<!std::is_same<T, V>::value>::type>
-    inline bool map(const V& from, T *to) { return rlookup(from, to); }
+    inline bool map(const V& from, T *to) const { return rlookup(from, to); }
 
 private:
     std::vector<std::pair<T, U>> mTable;
@@ -46,7 +46,7 @@
 }
 
 template<typename T, typename U>
-bool ALookup<T, U>::lookup(const T& from, U *to) {
+bool ALookup<T, U>::lookup(const T& from, U *to) const {
     for (auto elem : mTable) {
         if (elem.first == from) {
             *to = elem.second;
@@ -57,7 +57,7 @@
 }
 
 template<typename T, typename U>
-bool ALookup<T, U>::rlookup(const U& from, T *to) {
+bool ALookup<T, U>::rlookup(const U& from, T *to) const {
     for (auto elem : mTable) {
         if (elem.second == from) {
             *to = elem.first;
diff --git a/include/media/stagefright/foundation/ColorUtils.h b/include/media/stagefright/foundation/ColorUtils.h
index c4971bf..f01a210 100644
--- a/include/media/stagefright/foundation/ColorUtils.h
+++ b/include/media/stagefright/foundation/ColorUtils.h
@@ -129,6 +129,15 @@
     static status_t convertCodecColorAspectsToPlatformAspects(
             const ColorAspects &aspects, int32_t *range, int32_t *standard, int32_t *transfer);
 
+    // converts Other values to Unspecified
+    static void convertCodecColorAspectsToIsoAspects(
+            const ColorAspects &aspects,
+            int32_t *primaries, int32_t *transfer, int32_t *coeffs, bool *fullRange);
+    // converts unsupported values to Other
+    static void convertIsoColorAspectsToCodecAspects(
+            int32_t primaries, int32_t transfer, int32_t coeffs, bool fullRange,
+            ColorAspects &aspects);
+
     // updates Unspecified color aspects to their defaults based on the video size
     static void setDefaultCodecColorAspectsIfNeeded(
             ColorAspects &aspects, int32_t width, int32_t height);
diff --git a/media/libstagefright/MPEG4Extractor.cpp b/media/libstagefright/MPEG4Extractor.cpp
index d5a869d..322eab9 100644
--- a/media/libstagefright/MPEG4Extractor.cpp
+++ b/media/libstagefright/MPEG4Extractor.cpp
@@ -34,6 +34,7 @@
 #include <media/stagefright/foundation/ADebug.h>
 #include <media/stagefright/foundation/AMessage.h>
 #include <media/stagefright/foundation/AUtils.h>
+#include <media/stagefright/foundation/ColorUtils.h>
 #include <media/stagefright/MediaBuffer.h>
 #include <media/stagefright/MediaBufferGroup.h>
 #include <media/stagefright/MediaDefs.h>
@@ -2083,6 +2084,21 @@
             break;
         }
 
+        case FOURCC('c', 'o', 'l', 'r'):
+        {
+            *offset += chunk_size;
+            // this must be in a VisualSampleEntry box under the Sample Description Box ('stsd')
+            // ignore otherwise
+            if (depth >= 2 && mPath[depth - 2] == FOURCC('s', 't', 's', 'd')) {
+                status_t err = parseColorInfo(data_offset, chunk_data_size);
+                if (err != OK) {
+                    return err;
+                }
+            }
+
+            break;
+        }
+
         case FOURCC('t', 'i', 't', 'l'):
         case FOURCC('p', 'e', 'r', 'f'):
         case FOURCC('a', 'u', 't', 'h'):
@@ -2663,6 +2679,49 @@
     return OK;
 }
 
+status_t MPEG4Extractor::parseColorInfo(off64_t offset, size_t size) {
+    if (size < 4 || size == SIZE_MAX || mLastTrack == NULL) {
+        return ERROR_MALFORMED;
+    }
+
+    uint8_t *buffer = new (std::nothrow) uint8_t[size + 1];
+    if (buffer == NULL) {
+        return ERROR_MALFORMED;
+    }
+    if (mDataSource->readAt(offset, buffer, size) != (ssize_t)size) {
+        delete[] buffer;
+        buffer = NULL;
+
+        return ERROR_IO;
+    }
+
+    int32_t type = U32_AT(&buffer[0]);
+    if ((type == FOURCC('n', 'c', 'l', 'x') && size >= 11)
+            || (type == FOURCC('n', 'c', 'l', 'c' && size >= 10))) {
+        int32_t primaries = U16_AT(&buffer[4]);
+        int32_t transfer = U16_AT(&buffer[6]);
+        int32_t coeffs = U16_AT(&buffer[8]);
+        bool fullRange = (type == FOURCC('n', 'c', 'l', 'x')) && (buffer[10] & 128);
+
+        ColorAspects aspects;
+        ColorUtils::convertIsoColorAspectsToCodecAspects(
+                primaries, transfer, coeffs, fullRange, aspects);
+
+        // only store the first color specification
+        if (!mLastTrack->meta->hasData(kKeyColorPrimaries)) {
+            mLastTrack->meta->setInt32(kKeyColorPrimaries, aspects.mPrimaries);
+            mLastTrack->meta->setInt32(kKeyTransferFunction, aspects.mTransfer);
+            mLastTrack->meta->setInt32(kKeyColorMatrix, aspects.mMatrixCoeffs);
+            mLastTrack->meta->setInt32(kKeyColorRange, aspects.mRange);
+        }
+    }
+
+    delete[] buffer;
+    buffer = NULL;
+
+    return OK;
+}
+
 status_t MPEG4Extractor::parse3GPPMetaData(off64_t offset, size_t size, int depth) {
     if (size < 4 || size == SIZE_MAX) {
         return ERROR_MALFORMED;
diff --git a/media/libstagefright/MPEG4Writer.cpp b/media/libstagefright/MPEG4Writer.cpp
index ef0e17f..7c03886 100644
--- a/media/libstagefright/MPEG4Writer.cpp
+++ b/media/libstagefright/MPEG4Writer.cpp
@@ -30,6 +30,7 @@
 
 #include <media/stagefright/foundation/ADebug.h>
 #include <media/stagefright/foundation/AMessage.h>
+#include <media/stagefright/foundation/ColorUtils.h>
 #include <media/stagefright/MPEG4Writer.h>
 #include <media/stagefright/MediaBuffer.h>
 #include <media/stagefright/MetaData.h>
@@ -371,6 +372,7 @@
     void writeVmhdBox();
     void writeHdlrBox();
     void writeTkhdBox(uint32_t now);
+    void writeColrBox();
     void writeMp4aEsdsBox();
     void writeMp4vEsdsBox();
     void writeAudioFourCCBox();
@@ -2353,6 +2355,7 @@
         if (buffer->meta_data()->findInt32(kKeyIsCodecConfig, &isCodecConfig)
                 && isCodecConfig) {
             CHECK(!mGotAllCodecSpecificData);
+            mMeta = mSource->getFormat(); // get output format after format change
 
             if (mIsAvc) {
                 status_t err = makeAVCCodecSpecificData(
@@ -2960,9 +2963,32 @@
     }
 
     writePaspBox();
+    writeColrBox();
     mOwner->endBox();  // mp4v, s263 or avc1
 }
 
+void MPEG4Writer::Track::writeColrBox() {
+    ColorAspects aspects;
+    memset(&aspects, 0, sizeof(aspects));
+    // TRICKY: using | instead of || because we want to execute all findInt32-s
+    if (mMeta->findInt32(kKeyColorPrimaries, (int32_t*)&aspects.mPrimaries)
+            | mMeta->findInt32(kKeyTransferFunction, (int32_t*)&aspects.mTransfer)
+            | mMeta->findInt32(kKeyColorMatrix, (int32_t*)&aspects.mMatrixCoeffs)
+            | mMeta->findInt32(kKeyColorRange, (int32_t*)&aspects.mRange)) {
+        int32_t primaries, transfer, coeffs;
+        bool fullRange;
+        ColorUtils::convertCodecColorAspectsToIsoAspects(
+                aspects, &primaries, &transfer, &coeffs, &fullRange);
+        mOwner->beginBox("colr");
+        mOwner->writeFourcc("nclx");
+        mOwner->writeInt16(primaries);
+        mOwner->writeInt16(transfer);
+        mOwner->writeInt16(coeffs);
+        mOwner->writeInt8(fullRange ? 128 : 0);
+        mOwner->endBox(); // colr
+    }
+}
+
 void MPEG4Writer::Track::writeAudioFourCCBox() {
     const char *mime;
     bool success = mMeta->findCString(kKeyMIMEType, &mime);
diff --git a/media/libstagefright/MediaCodecSource.cpp b/media/libstagefright/MediaCodecSource.cpp
index 772a5c4..b5f7b12 100644
--- a/media/libstagefright/MediaCodecSource.cpp
+++ b/media/libstagefright/MediaCodecSource.cpp
@@ -781,6 +781,13 @@
 
             mAvailEncoderInputIndices.push_back(index);
             feedEncoderInputBuffers();
+        } else if (cbID == MediaCodec::CB_OUTPUT_FORMAT_CHANGED) {
+            status_t err = mEncoder->getOutputFormat(&mOutputFormat);
+            if (err != OK) {
+                signalEOS(err);
+                break;
+            }
+            convertMessageToMetaData(mOutputFormat, mMeta);
         } else if (cbID == MediaCodec::CB_OUTPUT_AVAILABLE) {
             int32_t index;
             size_t offset;
diff --git a/media/libstagefright/foundation/ColorUtils.cpp b/media/libstagefright/foundation/ColorUtils.cpp
index 30d5b45..99031ca 100644
--- a/media/libstagefright/foundation/ColorUtils.cpp
+++ b/media/libstagefright/foundation/ColorUtils.cpp
@@ -29,6 +29,7 @@
 typedef ColorAspects CA;
 typedef ColorUtils CU;
 
+const static
 ALookup<CU::ColorRange, CA::Range> sRanges{
     {
         { CU::kColorRangeLimited, CA::RangeLimited },
@@ -37,6 +38,7 @@
     }
 };
 
+const static
 ALookup<CU::ColorStandard, std::pair<CA::Primaries, CA::MatrixCoeffs>> sStandards {
     {
         { CU::kColorStandardUnspecified,    { CA::PrimariesUnspecified, CA::MatrixUnspecified } },
@@ -56,6 +58,7 @@
     }
 };
 
+const static
 ALookup<CU::ColorTransfer, CA::Transfer> sTransfers{
     {
         { CU::kColorTransferUnspecified,    CA::TransferUnspecified },
@@ -243,6 +246,97 @@
     }
 }
 
+const static
+ALookup<int32_t, ColorAspects::Primaries> sIsoPrimaries {
+    {
+        { 1, ColorAspects::PrimariesBT709_5 },
+        { 2, ColorAspects::PrimariesUnspecified },
+        { 4, ColorAspects::PrimariesBT470_6M },
+        { 5, ColorAspects::PrimariesBT601_6_625 },
+        { 6, ColorAspects::PrimariesBT601_6_525 /* main */},
+        { 7, ColorAspects::PrimariesBT601_6_525 },
+        // -- ITU T.832 201201 ends here
+        { 8, ColorAspects::PrimariesGenericFilm },
+        { 9, ColorAspects::PrimariesBT2020 },
+        { 10, ColorAspects::PrimariesOther /* XYZ */ },
+    }
+};
+
+const static
+ALookup<int32_t, ColorAspects::Transfer> sIsoTransfers {
+    {
+        { 1, ColorAspects::TransferSMPTE170M /* main */},
+        { 2, ColorAspects::TransferUnspecified },
+        { 4, ColorAspects::TransferGamma22 },
+        { 5, ColorAspects::TransferGamma28 },
+        { 6, ColorAspects::TransferSMPTE170M },
+        { 7, ColorAspects::TransferSMPTE240M },
+        { 8, ColorAspects::TransferLinear },
+        { 9, ColorAspects::TransferOther /* log 100:1 */ },
+        { 10, ColorAspects::TransferOther /* log 316:1 */ },
+        { 11, ColorAspects::TransferXvYCC },
+        { 12, ColorAspects::TransferBT1361 },
+        { 13, ColorAspects::TransferSRGB },
+        // -- ITU T.832 201201 ends here
+        { 14, ColorAspects::TransferSMPTE170M },
+        { 15, ColorAspects::TransferSMPTE170M },
+        { 16, ColorAspects::TransferST2084 },
+        { 17, ColorAspects::TransferST428 },
+    }
+};
+
+const static
+ALookup<int32_t, ColorAspects::MatrixCoeffs> sIsoMatrixCoeffs {
+    {
+        { 0, ColorAspects::MatrixOther },
+        { 1, ColorAspects::MatrixBT709_5 },
+        { 2, ColorAspects::MatrixUnspecified },
+        { 4, ColorAspects::MatrixBT470_6M },
+        { 6, ColorAspects::MatrixBT601_6 /* main */ },
+        { 5, ColorAspects::MatrixBT601_6 },
+        { 7, ColorAspects::MatrixSMPTE240M },
+        { 8, ColorAspects::MatrixOther /* YCgCo */ },
+        // -- ITU T.832 201201 ends here
+        { 9, ColorAspects::MatrixBT2020 },
+        { 10, ColorAspects::MatrixBT2020Constant },
+    }
+};
+
+// static
+void ColorUtils::convertCodecColorAspectsToIsoAspects(
+        const ColorAspects &aspects,
+        int32_t *primaries, int32_t *transfer, int32_t *coeffs, bool *fullRange) {
+    if (aspects.mPrimaries == ColorAspects::PrimariesOther ||
+            !sIsoPrimaries.map(aspects.mPrimaries, primaries)) {
+        CHECK(sIsoPrimaries.map(ColorAspects::PrimariesUnspecified, primaries));
+    }
+    if (aspects.mTransfer == ColorAspects::TransferOther ||
+            !sIsoTransfers.map(aspects.mTransfer, transfer)) {
+        CHECK(sIsoTransfers.map(ColorAspects::TransferUnspecified, transfer));
+    }
+    if (aspects.mMatrixCoeffs == ColorAspects::MatrixOther ||
+            !sIsoMatrixCoeffs.map(aspects.mMatrixCoeffs, coeffs)) {
+        CHECK(sIsoMatrixCoeffs.map(ColorAspects::MatrixUnspecified, coeffs));
+    }
+    *fullRange = aspects.mRange == ColorAspects::RangeFull;
+}
+
+// static
+void ColorUtils::convertIsoColorAspectsToCodecAspects(
+        int32_t primaries, int32_t transfer, int32_t coeffs, bool fullRange,
+        ColorAspects &aspects) {
+    if (!sIsoPrimaries.map(primaries, &aspects.mPrimaries)) {
+        aspects.mPrimaries = ColorAspects::PrimariesUnspecified;
+    }
+    if (!sIsoTransfers.map(transfer, &aspects.mTransfer)) {
+        aspects.mTransfer = ColorAspects::TransferUnspecified;
+    }
+    if (!sIsoMatrixCoeffs.map(coeffs, &aspects.mMatrixCoeffs)) {
+        aspects.mMatrixCoeffs = ColorAspects::MatrixUnspecified;
+    }
+    aspects.mRange = fullRange ? ColorAspects::RangeFull : ColorAspects::RangeLimited;
+}
+
 // static
 void ColorUtils::setDefaultCodecColorAspectsIfNeeded(
         ColorAspects &aspects, int32_t width, int32_t height) {
diff --git a/media/libstagefright/include/MPEG4Extractor.h b/media/libstagefright/include/MPEG4Extractor.h
index 3e8fb7c..18b14e1 100644
--- a/media/libstagefright/include/MPEG4Extractor.h
+++ b/media/libstagefright/include/MPEG4Extractor.h
@@ -110,6 +110,7 @@
     status_t readMetaData();
     status_t parseChunk(off64_t *offset, int depth);
     status_t parseITunesMetaData(off64_t offset, size_t size);
+    status_t parseColorInfo(off64_t offset, size_t size);
     status_t parse3GPPMetaData(off64_t offset, size_t size, int depth);
     void parseID3v2MetaData(off64_t offset);
     status_t parseQTMetaKey(off64_t data_offset, size_t data_size);