Add Stefan's io.py

scikits/image/io/io.py

@@ -0,0 +1,153 @@
+"""Read and write image files."""
+
+from __future__ import with_statement
+
+__all__ = ['Image','ImageCollection','imread']
+
+from glob import glob
+import os.path
+
+import numpy as N
+from numpy.testing import set_local_path, restore_path
+from scipy.misc.pilutil import imread
+
+set_local_path('../..')
+import supreme
+from supreme.lib import EXIF
+restore_path()
+
+class Image(N.ndarray):
+    """Image data with tags."""
+
+    tags = {'filename' : '',
+            'EXIF' : {},
+            'info' : {}}
+
+    def __new__(image_cls,arr,**kwargs):
+        """Set the image data and tags according to given parameters.
+
+        Input:
+        ------
+        `image_cls` : Image class specification
+            This is not normally specified by the user.
+        `arr` : ndarray
+            Image data.
+        ``**kwargs`` : Image tags as keywords
+            Specified in the form ``tag0=value``, ``tag1=value``.
+
+        """
+        x = N.asarray(arr).view(image_cls)
+        for tag,value in Image.tags.items():
+            setattr(x,tag,kwargs.get(tag,getattr(arr,tag,value)))
+        return x
+
+    def __array_finalize__(self, obj):
+        """Copy object tags."""
+        for tag,value in Image.tags.items():
+            setattr(self,tag,getattr(obj,tag,value))
+        return
+
+    def __reduce__(self):
+        object_state = list(N.ndarray.__reduce__(self))
+        subclass_state = {}
+        for tag in self.tags:
+            subclass_state[tag] = getattr(self,tag)
+        object_state[2] = (object_state[2],subclass_state)
+        return tuple(object_state)
+
+    def __setstate__(self,state):
+        nd_state,subclass_state = state
+        N.ndarray.__setstate__(self,nd_state)
+
+        for tag in subclass_state:
+            setattr(self,tag,subclass_state[tag])
+
+    @property
+    def exposure(self):
+        """Return exposure time based on EXIF tag."""
+        exposure = self.EXIF['EXIF ExposureTime'].values[0]
+        return exposure.num / float(exposure.den)
+
+class ImageCollection(object):
+    """Load and manage a collection of images."""
+
+    def __init__(self,file_pattern,conserve_memory=True,grey=False):
+        """Load image files.
+
+        Note that files are always stored in alphabetical order.
+
+        Input:
+        ------
+        file_pattern : string
+            Path and pattern of files to load, e.g. ``data/*.jpg``.
+        conserve_memory : bool
+            If True, never keep more than one in memory at a specific
+            time.  Otherwise, images will be cached once they are loaded.
+        grey : bool
+            If True, convert the input images to grey-scale.
+
+        Example:
+        --------
+        >>> from os.path import dirname, join
+        >>> data_dir = join(dirname(__file__),'tests')
+
+        >>> c = ImageCollection(data_dir + '/*.png')
+        >>> len(c)
+        3
+        >>> c[2].shape
+        (20, 20, 3)
+
+        """
+        self.files = sorted(glob(file_pattern))
+
+        if conserve_memory:
+            memory_slots = 1
+        else:
+            memory_slots = len(self.files)
+
+        self.conserve_memory = conserve_memory
+        self.grey = grey
+        self.data = N.empty(memory_slots,dtype=object)
+
+    def __getitem__(self,n,_cached=N.array(-1)):
+        """Return image n in the queue.
+
+        Loading is done on demand.
+
+        Input:
+        ------
+        n : int
+            Number of image required.
+
+        Output:
+        -------
+        img : array
+           Image #n in the collection.
+
+        """
+        idx = n % len(self.data)
+        if (_cached != n and self.conserve_memory) or (self.data[idx] is None):
+            image_data = imread(self.files[n],self.grey)
+
+            with file(self.files[n]) as f:
+                exif = EXIF.process_file(f)
+
+            self.data[idx] = Image(image_data,
+                                   filename=os.path.basename(self.files[n]),
+                                   EXIF=exif,info={})
+
+        _cached.flat = n
+
+        return self.data[idx]
+
+    def __iter__(self):
+        """Iterate over the images."""
+        for i in range(len(self)):
+            yield self[i]
+
+    def __len__(self):
+        """Number of images in collection."""
+        return len(self.files)
+
+    def __str__(self):
+        return str(self.files)