Fixing PNG writing for 2D

grafeo · Oct 8, 2015 · e1611ee · e1611ee
1 parent fc6df20
commit e1611ee
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/src/image.c b/src/image.c
@@ -53,6 +53,8 @@ Array* image_read(const char* filename){
   const char* ext = strrchr(filename, '.') + 1;
   if     (strcasecmp(ext, "png") == 0) return image_read_png(filename);
   else if(strcasecmp(ext, "jpg") == 0) return image_read_jpg(filename);
+  else if(strcasecmp(ext, "pgm") == 0) return image_read_pgm(filename);
+  else if(strcasecmp(ext, "ppm") == 0) return image_read_ppm(filename);
   return NULL;
 }
 Array* image_read_png(const char* filename){
@@ -189,6 +191,8 @@ void image_write(Array* array, const char* filename){
   const char* ext = strrchr(filename, '.') + 1;
   if     (strcasecmp(ext, "png") == 0) image_write_png(array,filename);
   else if(strcasecmp(ext, "jpg") == 0) image_write_jpg(array,filename);
+  else if(strcasecmp(ext, "pgm") == 0) image_write_pgm(array,filename);
+  else if(strcasecmp(ext, "ppm") == 0) image_write_ppm(array,filename);
 }
 void image_write_png(Array* array, const char* filename){
   /* create file */
@@ -212,7 +216,7 @@ void image_write_png(Array* array, const char* filename){
   if (setjmp(png_jmpbuf(png_ptr)))
     abort_("[write_png_file] Error during writing header");
 
-  int color_type = (array->size[2] == 3)?PNG_COLOR_TYPE_RGB:PNG_COLOR_TYPE_GRAY;
+  int color_type = (array->dim > 2 && array->size[2] == 3)?PNG_COLOR_TYPE_RGB:PNG_COLOR_TYPE_GRAY;
   png_set_IHDR(png_ptr, info_ptr, array->size[1], array->size[0],
                array->bitsize << 3, color_type, PNG_INTERLACE_NONE,
                PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
@@ -253,8 +257,8 @@ void image_write_jpg(Array* array, const char* filename){
   jpeg_stdio_dest(&cinfo, outfile);
   cinfo.image_width = array->size[1];
   cinfo.image_height = array->size[0];
-  cinfo.input_components = array->size[2];
-  cinfo.in_color_space = (array->size[2] == 3)?JCS_RGB:JCS_GRAYSCALE;
+  cinfo.input_components = (array->dim > 2)?array->size[2]:1;
+  cinfo.in_color_space = (array->dim > 2 && array->size[2] == 3)?JCS_RGB:JCS_GRAYSCALE;
   jpeg_set_defaults(&cinfo);
   jpeg_set_quality(&cinfo, 1, TRUE);
   jpeg_start_compress(&cinfo, TRUE);

diff --git a/tests/test_image.c b/tests/test_image.c
@@ -113,8 +113,9 @@ static void test_image_write_jpg(void ** state){
 
 static void test_image_write_png(void ** state){
   (void)state;
-  const char* outfile = "imagem.png";
+  const char* outfile      = "imagem.png";
   const char* outfile_gray = "imagemgray.png";
+  const char* outfile_2D   = "test_image_write_png_2D.png";
   // Remove the file if it exists
   remove(outfile);
   remove(outfile_gray);
@@ -131,8 +132,17 @@ static void test_image_write_png(void ** state){
   assert_int_equal(access(outfile, F_OK), 0);
   assert_int_equal(access(outfile_gray, F_OK), 0);
 
+  // Testing 2D image
+  uint8_t data[9] = {0,255,0,255,0,255,0,255,0};
+  uint32_t size[2] = {3,3};
+  Array* array_2D = array_from_data(data,2,size,GRAFEO_UINT8);
+  image_write_png(array_2D, outfile_2D);
+  assert_int_equal(access(outfile_2D, F_OK), 0);
+
+  // Memory
   array_free(input_image);
   array_free(input_image_gray);
+  array_free(array_2D);
 }
 
 static void test_image_read(void** state){