Index: src/_png.cpp
===================================================================
--- src/_png.cpp	(revision 7022)
+++ src/_png.cpp	(working copy)
@@ -208,38 +208,34 @@
 
   png_init_io(png_ptr, fp);
   png_set_sig_bytes(png_ptr, 8);
-
   png_read_info(png_ptr, info_ptr);
 
   png_uint_32 width = info_ptr->width;
   png_uint_32 height = info_ptr->height;
-  bool do_gray_conversion = (info_ptr->bit_depth < 8 &&
-                             info_ptr->color_type == PNG_COLOR_TYPE_GRAY);
 
   int bit_depth = info_ptr->bit_depth;
-  if (bit_depth == 16) {
-    png_set_strip_16(png_ptr);
-  } else if (bit_depth < 8) {
+  
+  if (bit_depth < 8) {
     png_set_packing(png_ptr);
   }
 
-  // convert misc color types to rgb for simplicity
-  if (info_ptr->color_type == PNG_COLOR_TYPE_GRAY ||
-      info_ptr->color_type == PNG_COLOR_TYPE_GRAY_ALPHA) {
-    png_set_gray_to_rgb(png_ptr);
-  } else if (info_ptr->color_type == PNG_COLOR_TYPE_PALETTE) {
+  // If sig bits are set, shift data
+  png_color_8p sig_bit;
+  if (png_get_sBIT(png_ptr, info_ptr, &sig_bit))
+    png_set_shift(png_ptr, sig_bit);
+
+  // Convert big endian to little
+  if (bit_depth == 16)
+    png_set_swap(png_ptr);
+
+  // Convert palletes to full RGB
+  if (info_ptr->color_type == PNG_COLOR_TYPE_PALETTE) {
     png_set_palette_to_rgb(png_ptr);
   }
 
   png_set_interlace_handling(png_ptr);
   png_read_update_info(png_ptr, info_ptr);
 
-  bool rgba = info_ptr->color_type == PNG_COLOR_TYPE_RGBA;
-  if ( (info_ptr->color_type != PNG_COLOR_TYPE_RGB) && !rgba) {
-    std::cerr << "Found color type " << (int)info_ptr->color_type  << std::endl;
-    throw Py::RuntimeError("_image_module::readpng: cannot handle color_type");
-  }
-
   /* read file */
   if (setjmp(png_jmpbuf(png_ptr)))
     throw Py::RuntimeError("_image_module::readpng: error during read_image");
@@ -255,35 +251,61 @@
   npy_intp dimensions[3];
   dimensions[0] = height;  //numrows
   dimensions[1] = width;   //numcols
-  dimensions[2] = 4;
+  dimensions[2] = 1;       //Dummy for now, set to 3/4 for RGB/RGBA
 
-  PyArrayObject *A = (PyArrayObject *) PyArray_SimpleNew(3, dimensions, PyArray_FLOAT);
+  PyArrayObject *A;
 
-  if (do_gray_conversion) {
-    float max_value = (float)((1L << bit_depth) - 1);
+  // Deal with each colour type and bit depth separately to support different return types
+  if ((bit_depth == 16) && (info_ptr->color_type == PNG_COLOR_TYPE_GRAY)) {
+    A = (PyArrayObject *) PyArray_SimpleNew(2, dimensions, PyArray_DOUBLE);
+	double max_value = (1 << (sig_bit->gray));
     for (png_uint_32 y = 0; y < height; y++) {
       png_byte* row = row_pointers[y];
-      for (png_uint_32 x = 0; x < width; x++) {
-        float value = row[x] / max_value;
+	  for (png_uint_32 x = 0; x < width; x++) {
         size_t offset = y*A->strides[0] + x*A->strides[1];
-        *(float*)(A->data + offset + 0*A->strides[2]) = value;
-        *(float*)(A->data + offset + 1*A->strides[2]) = value;
-        *(float*)(A->data + offset + 2*A->strides[2]) = value;
-        *(float*)(A->data + offset + 3*A->strides[2]) = 1.0f;
+        *(double*)(A->data + offset) = static_cast<double>(reinterpret_cast<png_uint_16*>(row)[x]) / max_value;
       }
     }
-  } else {
+  } else if (info_ptr->color_type == PNG_COLOR_TYPE_GRAY) {
+	A = (PyArrayObject *) PyArray_SimpleNew(2, dimensions, PyArray_FLOAT);
+	float max_value = ((1 << bit_depth) - 1);
+	for (png_uint_32 y = 0; y < height; y++) {
+	  png_byte* row = row_pointers[y];
+	  for (png_uint_32 x = 0; x < width; x++) {
+		size_t offset = y*A->strides[0] + x*A->strides[1];
+		*(float*)(A->data + offset) = static_cast<float>(row[x]) / max_value;
+	  }
+	}
+  } else if ((bit_depth == 8) && (info_ptr->color_type == PNG_COLOR_TYPE_RGBA)) {
+	dimensions[2] = 4;
+	A = (PyArrayObject *) PyArray_SimpleNew(3, dimensions, PyArray_FLOAT);
     for (png_uint_32 y = 0; y < height; y++) {
       png_byte* row = row_pointers[y];
       for (png_uint_32 x = 0; x < width; x++) {
-        png_byte* ptr = (rgba) ? &(row[x*4]) : &(row[x*3]);
+        png_byte* ptr = &(row[x*4]);
         size_t offset = y*A->strides[0] + x*A->strides[1];
-        *(float*)(A->data + offset + 0*A->strides[2]) = (float)(ptr[0]/255.0);
-        *(float*)(A->data + offset + 1*A->strides[2]) = (float)(ptr[1]/255.0);
-        *(float*)(A->data + offset + 2*A->strides[2]) = (float)(ptr[2]/255.0);
-        *(float*)(A->data + offset + 3*A->strides[2]) = rgba ? (float)(ptr[3]/255.0) : 1.0f;
+        *(float*)(A->data + offset + 0*A->strides[2]) = (float)(ptr[0]) / 255.;
+        *(float*)(A->data + offset + 1*A->strides[2]) = (float)(ptr[1]) / 255.;
+        *(float*)(A->data + offset + 2*A->strides[2]) = (float)(ptr[2]) / 255.;
+        *(float*)(A->data + offset + 3*A->strides[2]) = (float)(ptr[3]) / 255.;
       }
     }
+  } else if ((bit_depth == 8) && (info_ptr->color_type == PNG_COLOR_TYPE_RGB)) {
+    dimensions[2] = 3;
+	A = (PyArrayObject *) PyArray_SimpleNew(3, dimensions, PyArray_FLOAT);
+	for (png_uint_32 y = 0; y < height; y++) {
+	  png_byte* row = row_pointers[y];
+	  for (png_uint_32 x = 0; x < width; x++) {
+	    png_byte* ptr = &(row[x*3]);
+	    size_t offset = y*A->strides[0] + x*A->strides[1];
+		*(float*)(A->data + offset + 0*A->strides[2]) = (float)(ptr[0]) / 255.;
+		*(float*)(A->data + offset + 1*A->strides[2]) = (float)(ptr[1]) / 255.;
+		*(float*)(A->data + offset + 2*A->strides[2]) = (float)(ptr[2]) / 255.;
+	  }
+	}
+  } else {
+	std::cerr << "Found unsupported color type " << (int)info_ptr->color_type  << std::endl;
+    throw Py::RuntimeError("_image_module::readpng: cannot handle color_type");
   }
 
   //free the png memory