[Container] Add PNG support to new img container. (#380)

examples/BuddyLeNet/CMakeLists.txt

@@ -57,5 +57,5 @@ SET_TARGET_PROPERTIES(LENET PROPERTIES LINKER_LANGUAGE C)
 add_executable(buddy-lenet-run buddy-lenet-main.cpp)
 target_link_directories(buddy-lenet-run PRIVATE ${LLVM_LIBRARY_DIR})
 
-set(BUDDY_LENET_LIBS LENET mlir_c_runner_utils ${OpenCV_LIBS})
+set(BUDDY_LENET_LIBS LENET mlir_c_runner_utils ${OpenCV_LIBS} ${PNG_LIBRARY})
 target_link_libraries(buddy-lenet-run ${BUDDY_LENET_LIBS})

examples/BuddyLeNet/buddy-lenet-main.cpp

@@ -27,7 +27,7 @@
 #include <vector>
 
 constexpr size_t ParamsSize = 44426;
-const std::string ImgName = "8-16bit-565.bmp";
+const std::string ImgName = "0_1.png";
 
 /// Declare LeNet forward function.
 extern "C" void _mlir_ciface_forward(MemRef<float, 2> *output,

examples/BuddyMobileNetV3/CMakeLists.txt

@@ -71,5 +71,5 @@ SET_TARGET_PROPERTIES(MOBILENETV3 PROPERTIES LINKER_LANGUAGE C)
 add_executable(buddy-mobilenetv3-run buddy-mobilenetv3-main.cpp)
 target_link_directories(buddy-mobilenetv3-run PRIVATE ${LLVM_LIBRARY_DIR})
 
-set(BUDDY_MOBILENETV3_LIBS MOBILENETV3 mlir_c_runner_utils ${OpenCV_LIBS})
+set(BUDDY_MOBILENETV3_LIBS MOBILENETV3 mlir_c_runner_utils ${PNG_LIBRARY})
 target_link_libraries(buddy-mobilenetv3-run ${BUDDY_MOBILENETV3_LIBS})

examples/BuddyMobileNetV3/buddy-mobilenetv3-main.cpp

@@ -31,9 +31,9 @@ const std::string ImgName = "traffic-light-32bit.bmp";
 
 // Declare the mobilenet C interface.
 extern "C" void _mlir_ciface_forward(MemRef<float, 2> *output,
-                          MemRef<float, 1> *arg0,
-                          MemRef<long long, 1> *arg1,
-                          dip::Image<float, 4> *input);
+                                     MemRef<float, 1> *arg0,
+                                     MemRef<long long, 1> *arg1,
+                                     dip::Image<float, 4> *input);
 
 /// Print [Log] label in bold blue format.
 void printLogLabel() { std::cout << "\033[34;1m[Log] \033[0m"; }
@@ -55,7 +55,6 @@ void loadParameters(const std::string &floatParamPath,
   }
   floatParamFile.close();
 
-
   std::ifstream int64ParamFile(int64ParamPath, std::ios::in | std::ios::binary);
   if (!int64ParamFile.is_open()) {
     std::string errMsg = "Failed to open int64 param file: " +
@@ -94,8 +93,7 @@ void softmax(float *input, size_t size) {
 
 std::string getLabel(int idx) {
   std::string mobilenetDir = getenv("MOBILENETV3_EXAMPLE_PATH");
-  std::ifstream in(
-      mobilenetDir + "Labels.txt");
+  std::ifstream in(mobilenetDir + "Labels.txt");
   assert(in.is_open() && "Could not read the label file.");
   std::string label;
   for (int i = 0; i < idx; ++i)
@@ -111,15 +109,13 @@ int main() {
   std::cout << "\033[33;1m" << title << "\033[0m" << std::endl;
 
   // Define the sizes of the input and output tensors.
-  intptr_t sizesInput[4] = {1, 3, 224, 224};
   intptr_t sizesOutput[2] = {1, 1000};
 
   // Create input and output containers for the image and model output.
   std::string mobilenetDir = getenv("MOBILENETV3_EXAMPLE_PATH");
-  std::string imgPath = mobilenetDir + "/images/" + ImgName; 
-
+  std::string imgPath = mobilenetDir + "/images/" + ImgName;
   dip::Image<float, 4> input(imgPath, dip::DIP_RGB, true /* norm */);
-
+  
   MemRef<float, 2> output(sizesOutput);
 
   // Load model parameters from the specified file.
@@ -129,8 +125,9 @@ int main() {
   MemRef<long long, 1> ParamsContainerInt64({34});
   loadParameters(paramsDir, intDir, paramsContainerf32, ParamsContainerInt64);
   // Call the forward function of the model.
-  _mlir_ciface_forward(&output, &paramsContainerf32, &ParamsContainerInt64, &input);
-
+  _mlir_ciface_forward(&output, &paramsContainerf32, &ParamsContainerInt64,
+                       &input);
+
   auto out = output.getData();
   softmax(out, 1000);
   // Find the classification and print the result.
@@ -143,7 +140,7 @@ int main() {
     }
   }
   std::cout << "Classification Index: " << maxIdx << std::endl;
-  std::cout << "Classification: " << getLabel(maxIdx+1) << std::endl;
+  std::cout << "Classification: " << getLabel(maxIdx + 1) << std::endl;
   std::cout << "Probability: " << maxVal << std::endl;
 
   return 0;

frontend/Interfaces/buddy/DIP/ImgContainer.h

@@ -25,13 +25,20 @@
 #include <cstring>
 #include <fstream>
 #include <memory>
+#include <png.h>
 
 namespace dip {
 enum ImageModes {
   DIP_GRAYSCALE = 0,
   DIP_RGB = 1,
 };
 
+inline bool isBigEndian() {
+  int num = 1;
+  char *ptr = (char *)&num;
+  return (*ptr == 0);
+}
+
 template <typename T, size_t N> class Image : public MemRef<T, N> {
 public:
   // Constructor initializes the image by loading from a file.
@@ -46,6 +53,8 @@ template <typename T, size_t N> class Image : public MemRef<T, N> {
     switch (this->imageFormat) {
     case ImageFormat::BMP:
       return "BMP";
+    case ImageFormat::PNG:
+      return "PNG";
     default:
       return "Unsupported format";
     }
@@ -62,6 +71,7 @@ template <typename T, size_t N> class Image : public MemRef<T, N> {
   enum class ImageFormat {
     ERROR, // Represents an error or unsupported format.
     BMP,   // BMP file format.
+    PNG,   // PNG file format.
   } imageFormat;
   // Mode of the image (e.g., DIP_GRAYSCALE, DIP_RGB).
   ImageModes imageMode;
@@ -77,6 +87,8 @@ template <typename T, size_t N> class Image : public MemRef<T, N> {
   void determineFormat(const std::vector<uint8_t> &fileData);
   // Decodes a BMP image from raw file data.
   bool decodeBMP(const std::vector<uint8_t> &fileData);
+  // Decodes a PNG image from raw file data.
+  bool decodePNG(const std::vector<uint8_t> &fileData);
 };
 
 // Image Container Constructor
@@ -117,6 +129,12 @@ Image<T, N>::Image(std::string filePath, ImageModes mode, bool norm)
       this->imageFormat = ImageFormat::ERROR;
       throw std::runtime_error("Failed to decode BMP file from " + filePath);
     };
+  } else if (this->imageFormat == ImageFormat::PNG) {
+    bool success = decodePNG(fileData);
+    if (!success) {
+      this->imageFormat = ImageFormat::ERROR;
+      throw std::runtime_error("Failed to decode PNG file from " + filePath);
+    };
   } else {
     throw std::runtime_error("Unsupported image file format.");
   }
@@ -125,8 +143,13 @@ Image<T, N>::Image(std::string filePath, ImageModes mode, bool norm)
 // Determines the image format by inspecting the header of the file data.
 template <typename T, std::size_t N>
 void Image<T, N>::determineFormat(const std::vector<uint8_t> &fileData) {
+  std::array<unsigned char, 8> pngHeader = {0x89, 0x50, 0x4E, 0x47,
+                                            0x0D, 0x0A, 0x1A, 0x0A};
   if (fileData.size() > 2 && fileData[0] == 'B' && fileData[1] == 'M') {
     this->imageFormat = ImageFormat::BMP;
+  } else if (fileData.size() > 7 &&
+             std::memcmp(fileData.data(), pngHeader.data(), 8) == 0) {
+    this->imageFormat = ImageFormat::PNG;
   } else {
     this->imageFormat = ImageFormat::ERROR;
   }
@@ -390,6 +413,199 @@ bool Image<T, N>::decodeBMP(const std::vector<uint8_t> &fileData) {
   return true;
 }
 
+// PNG Image File Decoder
+template <typename T, std::size_t N>
+bool Image<T, N>::decodePNG(const std::vector<uint8_t> &fileData) {
+  // Check if the provided data is large enough to contain a minimal PNG header
+  // (33 bytes).
+  if (fileData.size() < 33) {
+    throw std::runtime_error("Invalid PNG File: too small to contain header");
+  }
+
+  // Extract image information from PNG header. Convert Big-Endian to
+  // Little-Endian.
+  this->width = (fileData[16] << 24) | (fileData[17] << 16) |
+                (fileData[18] << 8) | fileData[19];
+  this->height = (fileData[20] << 24) | (fileData[21] << 16) |
+                 (fileData[22] << 8) | fileData[23];
+  this->bitDepth = *reinterpret_cast<const uint8_t *>(&fileData[24]);
+  int colorType = *reinterpret_cast<const uint8_t *>(&fileData[25]);
+  uint8_t interlace = *reinterpret_cast<const uint8_t *>(&fileData[28]);
+
+  // Currently, only the NCHW format with 4 dimensions is supported.
+  if (N == 4) {
+    // use libpng to read png image. Initialize libpng parameters
+    png_structp png_ptr =
+        png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
+    if (!png_ptr) {
+      std::cerr << "png_ptr creation failed" << std::endl;
+      return false;
+    }
+
+    png_infop info_ptr = png_create_info_struct(png_ptr);
+    if (!info_ptr) {
+      std::cerr << "png_infop creation failed" << std::endl;
+      return false;
+    }
+
+    // Set jump point for error handling
+    if (setjmp(png_jmpbuf(png_ptr))) {
+      std::cerr << "error during PNG reading" << std::endl;
+      // close PNG reading and free memory
+      png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
+      return false;
+    }
+
+    // copy filedata. Read image data from memory.
+    std::vector<uint8_t> dataCopy = fileData;
+    png_set_read_fn(
+        png_ptr, &dataCopy,
+        [](png_structp png_ptr, png_bytep data, png_size_t length) {
+          std::vector<uint8_t> *fileData =
+              static_cast<std::vector<uint8_t> *>(png_get_io_ptr(png_ptr));
+          if (fileData->size() < length) {
+            png_error(png_ptr, "Read error from memory");
+          }
+          std::copy(fileData->begin(), fileData->begin() + length, data);
+          fileData->erase(fileData->begin(), fileData->begin() + length);
+        });
+
+    png_read_info(png_ptr, info_ptr);
+
+    // Convert big or little Endian and convert 16 bits to 8 bits
+    if (this->bitDepth == 16)
+      png_set_strip_16(png_ptr);
+    else if (!isBigEndian())
+      png_set_swap(png_ptr);
+
+    // Remove alpha channel
+    if (colorType & PNG_COLOR_MASK_ALPHA)
+      png_set_strip_alpha(png_ptr);
+
+    // Convert palette to rgb
+    if (colorType == PNG_COLOR_TYPE_PALETTE)
+      png_set_palette_to_rgb(png_ptr);
+
+    // Convert low depth grayscale to 8-bit grayscale
+    if ((colorType & PNG_COLOR_MASK_COLOR) == 0 && this->bitDepth < 8)
+#if (PNG_LIBPNG_VER_MAJOR * 10000 + PNG_LIBPNG_VER_MINOR * 100 +               \
+         PNG_LIBPNG_VER_RELEASE >=                                             \
+     10209) ||                                                                 \
+    (PNG_LIBPNG_VER_MAJOR == 1 && PNG_LIBPNG_VER_MINOR == 0 &&                 \
+     PNG_LIBPNG_VER_RELEASE >= 18)
+      png_set_expand_gray_1_2_4_to_8(png_ptr);
+#else
+      png_set_gray_1_2_4_to_8(png_ptr);
+#endif
+
+    // Processing interleaved PNG images
+    if (interlace)
+      png_set_interlace_handling(png_ptr);
+
+    if (this->imageMode == ImageModes::DIP_GRAYSCALE) {
+      // TODO: Add batch setting.
+      this->sizes[0] = 1;
+      this->sizes[1] = 1;
+      this->sizes[2] = this->height;
+      this->sizes[3] = this->width;
+      this->setStrides();
+      size_t size = this->product(this->sizes);
+      this->allocated = (T *)malloc(sizeof(T) * size);
+      this->aligned = this->allocated;
+
+      // RGB->Gray
+      if ((colorType & PNG_COLOR_MASK_COLOR) ||
+          (colorType == PNG_COLOR_TYPE_PALETTE))
+        png_set_rgb_to_gray(png_ptr, 1, 0.299, 0.587);
+
+      // Update reading setting
+      png_read_update_info(png_ptr, info_ptr);
+
+      // Allocate memory for libpng to read images
+      std::vector<uint8_t> imgData(this->height * this->width);
+      std::vector<uint8_t *> row_pointers(this->height);
+      for (size_t y = 0; y < this->height; ++y) {
+        row_pointers[y] = imgData.data() + y * this->width;
+      }
+
+      // Reading image
+      png_read_image(png_ptr, row_pointers.data());
+
+      // Fullfill data to memref container.
+      for (size_t i = 0; i < this->height; i++)
+        for (size_t j = 0; j < this->width; j++) {
+          size_t memrefIndex = i * this->width + j;
+          this->aligned[memrefIndex] =
+              this->isNorm ? static_cast<T>(imgData[memrefIndex]) / 255
+                           : static_cast<T>(imgData[memrefIndex]);
+          ;
+        }
+    } else if (this->imageMode == ImageModes::DIP_RGB) {
+      // TODO: Add batch setting.
+      this->sizes[0] = 1;
+      this->sizes[1] = 3;
+      this->sizes[2] = this->height;
+      this->sizes[3] = this->width;
+      this->setStrides();
+      size_t size = this->product(this->sizes);
+      std::cout << "size" << size << std::endl;
+      this->allocated = (T *)malloc(sizeof(T) * size);
+      this->aligned = this->allocated;
+      size_t colorStride = this->height * this->width;
+
+      // Gray->RGB
+      if (colorType & PNG_COLOR_TYPE_GRAY)
+        png_set_gray_to_rgb(png_ptr);
+
+      // Update reading setting
+      png_read_update_info(png_ptr, info_ptr);
+
+      // Allocate memory for libpng to read images
+      std::vector<uint8_t> imgData(this->height * this->width * 3);
+      std::vector<uint8_t *> row_pointers(this->height);
+      for (size_t y = 0; y < this->height; ++y) {
+        row_pointers[y] = imgData.data() + y * this->width * 3;
+      }
+
+      // Reading image
+      png_read_image(png_ptr, row_pointers.data());
+
+      // Separate pixel data by channel
+      size_t memrefIndex = 0;
+      for (size_t i = 0; i < this->height; i++)
+        for (size_t j = 0; j < this->width; j++) {
+          // Locate the current pixel.
+          size_t pixelIndex = ((i * width) + j) * 3;
+          // Extract the red, green, and blue value from the current pixel.
+          int redPixel =
+              *reinterpret_cast<const uint8_t *>(&imgData[pixelIndex]);
+          int greenPixel =
+              *reinterpret_cast<const uint8_t *>(&imgData[pixelIndex + 1]);
+          int bluePixel =
+              *reinterpret_cast<const uint8_t *>(&imgData[pixelIndex + 2]);
+          // Store the values into memref container as RGB order.
+          this->aligned[memrefIndex] = this->isNorm
+                                           ? static_cast<T>(redPixel) / 255
+                                           : static_cast<T>(redPixel);
+          this->aligned[memrefIndex + colorStride] =
+              this->isNorm ? static_cast<T>(greenPixel) / 255
+                           : static_cast<T>(greenPixel);
+          this->aligned[memrefIndex + 2 * colorStride] =
+              this->isNorm ? static_cast<T>(bluePixel) / 255
+                           : static_cast<T>(bluePixel);
+          memrefIndex++;
+        }
+    }
+
+    // close PNG reading and free memory
+    png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
+  } else {
+    std::cerr << "Unsupported: " << N << " dimension layout." << std::endl;
+    return false;
+  }
+  return true;
+}
+
 } // namespace dip
 
 #endif // FRONTEND_INTERFACES_BUDDY_DIP_IMGCONTAINER

tests/CMakeLists.txt

@@ -13,7 +13,6 @@ set(BUDDY_TEST_DEPENDS
   buddy-translate
   buddy-container-test
   buddy-audio-container-test
-  buddy-new-image-container-test
   buddy-text-container-test
   mlir-cpu-runner
   )
@@ -22,6 +21,10 @@ if(BUDDY_ENABLE_OPENCV)
   list(APPEND BUDDY_TEST_DEPENDS buddy-image-container-test)
 endif()
 
+if(BUDDY_MLIR_ENABLE_DIP_LIB)
+  list(APPEND BUDDY_TEST_DEPENDS buddy-new-image-container-test)
+endif()
+
 add_lit_testsuite(check-tests "Running the buddy regression tests..."
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${BUDDY_TEST_DEPENDS}