Allow record to set input color format as a command line option

o Only YUV420SP and YUV420P are supported

Change-Id: I425e8b29d9dca8c05066c1cd777f8bdede52964a
diff --git a/cmds/stagefright/record.cpp b/cmds/stagefright/record.cpp
index c424281..b718299 100644
--- a/cmds/stagefright/record.cpp
+++ b/cmds/stagefright/record.cpp
@@ -36,24 +36,30 @@
 static const int32_t kIFramesIntervalSec = 1;
 static const int32_t kVideoBitRate = 512 * 1024;
 static const int32_t kAudioBitRate = 12200;
-static const int32_t kColorFormat = OMX_COLOR_FormatYUV420SemiPlanar;
 static const int64_t kDurationUs = 10000000LL;  // 10 seconds
 
 #if 1
 class DummySource : public MediaSource {
 
 public:
-    DummySource(int width, int height)
+    DummySource(int width, int height, int colorFormat)
         : mWidth(width),
           mHeight(height),
+          mColorFormat(colorFormat),
           mSize((width * height * 3) / 2) {
         mGroup.add_buffer(new MediaBuffer(mSize));
+
+        // Check the color format to make sure
+        // that the buffer size mSize it set correctly above.
+        CHECK(colorFormat == OMX_COLOR_FormatYUV420SemiPlanar ||
+              colorFormat == OMX_COLOR_FormatYUV420Planar);
     }
 
     virtual sp<MetaData> getFormat() {
         sp<MetaData> meta = new MetaData;
         meta->setInt32(kKeyWidth, mWidth);
         meta->setInt32(kKeyHeight, mHeight);
+        meta->setInt32(kKeyColorFormat, mColorFormat);
         meta->setCString(kKeyMIMEType, MEDIA_MIMETYPE_VIDEO_RAW);
 
         return meta;
@@ -100,6 +106,7 @@
 private:
     MediaBufferGroup mGroup;
     int mWidth, mHeight;
+    int mColorFormat;
     size_t mSize;
     int64_t mNumFramesOutput;;
 
@@ -139,20 +146,47 @@
     return source;
 }
 
+enum {
+    kYUV420SP = 0,
+    kYUV420P  = 1,
+};
+
+// returns -1 if mapping of the given color is unsuccessful
+// returns an omx color enum value otherwise
+static int translateColorToOmxEnumValue(int color) {
+    switch (color) {
+        case kYUV420SP:
+            return OMX_COLOR_FormatYUV420SemiPlanar;
+        case kYUV420P:
+            return OMX_COLOR_FormatYUV420Planar;
+        default:
+            fprintf(stderr, "Unsupported color: %d\n", color);
+            return -1;
+    }
+}
+
 int main(int argc, char **argv) {
     android::ProcessState::self()->startThreadPool();
 
     DataSource::RegisterDefaultSniffers();
 
 #if 1
-    if (argc != 2) {
-        fprintf(stderr, "usage: %s filename\n", argv[0]);
+    if (argc != 3) {
+        fprintf(stderr, "usage: %s <filename> <input_color_format>\n", argv[0]);
+        fprintf(stderr, "       <input_color_format>:  0 (YUV420SP) or 1 (YUV420P)\n");
         return 1;
     }
 
+    int colorFormat = translateColorToOmxEnumValue(atoi(argv[2]));
+    if (colorFormat == -1) {
+        fprintf(stderr, "input color format must be 0 (YUV420SP) or 1 (YUV420P)\n");
+        return 1;
+    }
     OMXClient client;
     CHECK_EQ(client.connect(), OK);
 
+    status_t err = OK;
+
 #if 0
     sp<MediaSource> source = createSource(argv[1]);
 
@@ -173,7 +207,7 @@
 #else
     int width = 720;
     int height = 480;
-    sp<MediaSource> decoder = new DummySource(width, height);
+    sp<MediaSource> decoder = new DummySource(width, height, colorFormat);
 #endif
 
     sp<MetaData> enc_meta = new MetaData;
@@ -187,7 +221,7 @@
     enc_meta->setInt32(kKeyStride, width);
     enc_meta->setInt32(kKeySliceHeight, height);
     enc_meta->setInt32(kKeyIFramesInterval, kIFramesIntervalSec);
-    enc_meta->setInt32(kKeyColorFormat, kColorFormat);
+    enc_meta->setInt32(kKeyColorFormat, colorFormat);
 
     sp<MediaSource> encoder =
         OMXCodec::Create(
@@ -197,14 +231,14 @@
     sp<MPEG4Writer> writer = new MPEG4Writer("/sdcard/output.mp4");
     writer->addSource(encoder);
     writer->setMaxFileDuration(kDurationUs);
-    writer->start();
+    CHECK_EQ(OK, writer->start());
     while (!writer->reachedEOS()) {
         fprintf(stderr, ".");
         usleep(100000);
     }
-    writer->stop();
+    err = writer->stop();
 #else
-    encoder->start();
+    CHECK_EQ(OK, encoder->start());
 
     MediaBuffer *buffer;
     while (encoder->read(&buffer) == OK) {
@@ -222,7 +256,7 @@
         buffer = NULL;
     }
 
-    encoder->stop();
+    err = encoder->stop();
 #endif
 
     printf("$\n");
@@ -247,12 +281,16 @@
         buffer = NULL;
     }
 
-    source->stop();
+    err = source->stop();
 
     delete source;
     source = NULL;
 #endif
 
+    if (err != OK && err != ERROR_END_OF_STREAM) {
+        fprintf(stderr, "record failed: %d\n", err);
+        return 1;
+    }
     return 0;
 }
 #else