Step 2: bind function to Python and use it instead of imghdr

bob/io/image/__init__.py

@@ -9,6 +9,8 @@ from . import _library
 from . import version
 from .version import module as __version__
 
+from ._library import *
+
 import os
 
 def get_config():
@@ -17,13 +19,6 @@ def get_config():
   import bob.extension
   return bob.extension.get_config(__name__, version.externals)
 
-# fix imghdr's jpeg detection to use the first two bytes (according to https://en.wikipedia.org/wiki/List_of_file_signatures)
-import imghdr
-def _test_jpeg(h, f):
-  if h.startswith('\xff\xd8'):
-    return "jpeg"
-imghdr.tests.append(_test_jpeg)
-
 def load(filename, extension=None):
   """load(filename) -> image
 
@@ -54,10 +49,7 @@ def load(filename, extension=None):
     f = bob.io.base.File(filename, 'r')
   else:
     if extension == 'auto':
-      extension = imghdr.what(filename)
-      if extension is None:
-        raise IOError("Could not detect the image type of file %s" % filename)
-      extension = "." + extension
+      extension = get_correct_image_extension(filename)
     f = bob.io.base.File(filename, 'r', extension)
 
   return f.read()

bob/io/image/main.cpp

@@ -187,6 +187,29 @@ BOB_CATCH_FUNCTION("_test_io", 0)
 }
 
 
+static auto s_image_extension = bob::extension::FunctionDoc(
+  "get_correct_image_extension",
+  "Estimates the image type and return a corresponding extension based on file content",
+  "This function loads the first bytes of the given image, and matches it with known magic numbers of image files. "
+  "If a match is found, it returns the corresponding image extension (including the leading ``'.'`` that can be used, e.g., in :py:func:`bob.io.image.load`."
+)
+.add_prototype("image_name", "extension")
+.add_parameter("image_name", "str", "The name (including path) of the image to check")
+.add_return("extension", "str", "The extension of the image based on the file content")
+;
+static PyObject* image_extension(PyObject*, PyObject *args, PyObject* kwds) {
+BOB_TRY
+  static char** kwlist = s_image_extension.kwlist();
+
+  const char* image_name;
+  if (!PyArg_ParseTupleAndKeywords(args, kwds, "s", kwlist, &image_name)) return 0;
+
+  return Py_BuildValue("s", bob::io::image::get_correct_image_extension(image_name).c_str());
+
+BOB_CATCH_FUNCTION("get_correct_image_extension", 0)
+}
+
+
 static PyMethodDef module_methods[] = {
   {
     s_test_io.name(),
@@ -194,6 +217,12 @@ static PyMethodDef module_methods[] = {
     METH_VARARGS|METH_KEYWORDS,
     s_test_io.doc(),
   },
+  {
+    s_image_extension.name(),
+    (PyCFunction)image_extension,
+    METH_VARARGS|METH_KEYWORDS,
+    s_image_extension.doc(),
+  },
   {0}  /* Sentinel */
 };