PixelFrame helper methods

Summary: Add some helper methods and test for jpg loading Differential Revision: D58492082 fbshipit-source-id: 47f94ed66bd8b90b1f5b657397edf1db7b2c13b8
facebookresearch · Jun 13, 2024 · afe7d90 · afe7d90
1 parent 432fb40
commit afe7d90
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 4 deletions.
diff --git a/tools/vrsplayer/FramePlayer.cpp b/tools/vrsplayer/FramePlayer.cpp
@@ -265,7 +265,7 @@ void FramePlayer::imageJobsThreadActivity() {
     }
     shared_ptr<PixelFrame>& frame = *job;
     bool frameValid = false;
-    vrs::ImageFormat imageFormat = frame->getSpec().getImageFormat();
+    vrs::ImageFormat imageFormat = frame->getImageFormat();
     if (imageFormat == vrs::ImageFormat::RAW) {
       frameValid = (frame != nullptr);
     } else if (imageFormat == vrs::ImageFormat::VIDEO) {

diff --git a/vrs/utils/PixelFrame.h b/vrs/utils/PixelFrame.h
@@ -113,6 +113,9 @@ class PixelFrame {
   const ImageContentBlockSpec& getSpec() const {
     return imageSpec_;
   }
+  ImageFormat getImageFormat() const {
+    return imageSpec_.getImageFormat();
+  }
   PixelFormat getPixelFormat() const {
     return imageSpec_.getPixelFormat();
   }
@@ -122,9 +125,24 @@ class PixelFrame {
   uint32_t getHeight() const {
     return imageSpec_.getHeight();
   }
-  size_t getStride() const {
+  uint32_t getStride() const {
     return imageSpec_.getStride();
   }
+  uint32_t getDefaultStride() const {
+    return imageSpec_.getDefaultStride();
+  }
+  uint32_t getPlaneStride(uint32_t planeIndex) const {
+    return imageSpec_.getPlaneStride(planeIndex);
+  }
+  uint32_t getPlaneHeight(uint32_t planeIndex) const {
+    return imageSpec_.getPlaneHeight(planeIndex);
+  }
+  uint8_t getChannelCountPerPixel() const {
+    return imageSpec_.getChannelCountPerPixel();
+  }
+  size_t getBytesPerPixel() const {
+    return imageSpec_.getBytesPerPixel();
+  }
 
   vector<uint8_t>& getBuffer() {
     return frameBytes_;

diff --git a/vrs/utils/PixelFrameJpeg.cpp b/vrs/utils/PixelFrameJpeg.cpp
@@ -43,7 +43,6 @@ bool PixelFrame::readJpegFrame(RecordReader* reader, uint32_t sizeBytes) {
 static bool
 readJpegFrameHelper(PixelFrame& frame, struct jpeg_decompress_struct& cinfo, bool decodePixels) {
   jpeg_read_header(&cinfo, TRUE);
-  jpeg_start_decompress(&cinfo);
   if (cinfo.num_components == 1) {
     cinfo.out_color_space = JCS_GRAYSCALE;
     frame.init(PixelFormat::GREY8, cinfo.image_width, cinfo.image_height);
@@ -53,13 +52,14 @@ readJpegFrameHelper(PixelFrame& frame, struct jpeg_decompress_struct& cinfo, boo
   }
   if (decodePixels) {
     // decompress row by row
+    jpeg_start_decompress(&cinfo);
     uint8_t* rowPtr = frame.wdata();
     while (cinfo.output_scanline < cinfo.output_height) {
       jpeg_read_scanlines(&cinfo, &rowPtr, 1);
       rowPtr += frame.getSpec().getStride();
     }
+    jpeg_finish_decompress(&cinfo);
   }
-  jpeg_finish_decompress(&cinfo);
   jpeg_destroy_decompress(&cinfo);
   return true;
 }

diff --git a/vrs/utils/test/PixelFrameTest.cpp b/vrs/utils/test/PixelFrameTest.cpp
@@ -37,6 +37,7 @@ using coretech::getTestDataDir;
 
 struct PixelFrameTest : testing::Test {
   string kTestFilePath = os::pathJoin(getTestDataDir(), "VRS_Files/sample_raw_pixel_formats.vrs");
+  string kJpegTestFilePath = os::pathJoin(getTestDataDir(), "computer_vision/eileen_128x128.jpg");
 };
 
 class ImagePlayer : public RecordFormatStreamPlayer {
@@ -179,3 +180,20 @@ TEST_F(PixelFrameTest, colorTableTest) {
   EXPECT_TRUE(equal(colors[0], RGBColor()));
   EXPECT_TRUE(equal(colors[0xffff], RGBColor(255, 255, 255)));
 }
+
+#if IS_VRS_FB_INTERNAL()
+TEST_F(PixelFrameTest, jpegTest) {
+  PixelFrame frame;
+  ASSERT_TRUE(frame.readJpegFrameFromFile(kJpegTestFilePath, true));
+  EXPECT_EQ(frame.size(), 49152);
+  EXPECT_EQ(frame.data<uint64_t>()[0], 4198000490383812399);
+  EXPECT_EQ(frame.data<uint64_t>()[2500], 8221608522182776091);
+  EXPECT_EQ(frame.data<uint64_t>()[49152 / 8 - 1], 1099511693312);
+  PixelFrame frame2;
+  ASSERT_TRUE(frame2.readJpegFrameFromFile(kJpegTestFilePath, false));
+  EXPECT_EQ(frame2.size(), frame.size());
+  ASSERT_EQ(frame.getSpec(), frame2.getSpec());
+  frame.blankFrame();
+  EXPECT_EQ(frame.getBuffer(), frame2.getBuffer());
+}
+#endif