From 8da0ea52c01e5551aa62828507739ebd2846d215 Mon Sep 17 00:00:00 2001
From: Stefan van der Walt <stefan@sun.ac.za>
Date: Sun, 25 Sep 2011 15:15:07 -0700
Subject: [PATCH] ENH: Minor clean ups of data loader.

---
 scikits/image/data/__init__.py        | 44 +++++++++++++++++++--------
 scikits/image/data/tests/test_data.py | 33 ++++++++++----------
 2 files changed, 48 insertions(+), 29 deletions(-)

diff --git a/scikits/image/data/__init__.py b/scikits/image/data/__init__.py
index 095806b6..728b20dd 100644
--- a/scikits/image/data/__init__.py
+++ b/scikits/image/data/__init__.py
@@ -1,24 +1,42 @@
-"""Convinience functions to get sample data"""
+"""Convenience functions to load sample data.
 
-import os
+"""
+
+import os as _os
 
 from ..io import imread
 from ...image import data_dir
 
-def camera():
-    """Example gray "camera" image, often used for segmentation
-    and denoising examples."""
+def load(f):
+    """Load an image file located in the data directory.
 
-    return imread(os.path.join(data_dir, "camera.png"))
+    Parameters
+    ----------
+    f : string
+        File name.
+
+    Returns
+    -------
+    img : ndarray
+        Image loaded from scikits.image.data_dir.
+    """
+    return imread(_os.path.join(data_dir, f))
+
+def camera():
+    """Gray "camera" image, often used for segmentation
+    and denoising examples.
+
+    """
+    return load("camera.png")
 
 def lena():
-    """Example "Lena" image. """
-    return imread(os.path.join(data_dir, "lena.png"))
+    """Colour "Lena" image.
+
+    """
+    return load("lena.png")
 
 def checkerboard():
-    """Checkerboard image"""
-    return imread(os.path.join(data_dir, "chessboard_RGB.png"))
+    """Checkerboard image.
 
-def checkerboard_gray():
-    """Checkerboard image, only gray channel"""
-    return imread(os.path.join(data_dir, "chessboard_GRAY.png"))
+    """
+    return load("chessboard_RGB.png")
diff --git a/scikits/image/data/tests/test_data.py b/scikits/image/data/tests/test_data.py
index 77b8ddaa..b31c8851 100644
--- a/scikits/image/data/tests/test_data.py
+++ b/scikits/image/data/tests/test_data.py
@@ -1,22 +1,23 @@
-import unittest
-
 import scikits.image.data as data
+from numpy.testing import assert_equal, assert_array_equal
+import numpy as np
 
-class TestData(unittest.TestCase):
-    def test_lena(self):
-        """ Test that "Lena" image can be loaded. """
-        data.lena()
+def test_lena():
+    """ Test that "Lena" image can be loaded. """
+    lena = data.lena()
+    assert_equal(lena.shape, (512, 512, 3))
 
-    def test_camera(self):
-        """ Test that "camera" image can be loaded. """
-        data.camera()
+def test_camera():
+    """ Test that "camera" image can be loaded. """
+    cameraman = data.camera()
+    assert_equal(cameraman.ndim, 2)
 
-    def test_checkerboard(self):
-        """ Test that checkerboard image can be loaded. """
-        data.checkerboard()
-
-    def test_checkerboard_gray(self):
-        """ Test that checkerboard grayscale image can be loaded. """
-        data.checkerboard_gray()
+def test_checkerboard():
+    """ Test that checkerboard image can be loaded. """
+    checkerboard = data.checkerboard()
+    assert_equal(checkerboard.dtype, np.uint8)
 
+if __name__ == "__main__":
+    from numpy.testing import run_module_suite
+    run_module_suite()