From 57940efae0dbe2198e5fad17ec3ed96514d764d8 Mon Sep 17 00:00:00 2001
From: jim <jim@rybarski.com>
Date: Thu, 13 Nov 2014 14:14:49 -0600
Subject: [PATCH] got image set iteration working

---
 nd2reader/__init__.py         | 26 ++++++++++++++++++++++----
 nd2reader/model/__init__.py   | 18 ++++++++++++++----
 nd2reader/service/__init__.py |  4 ++++
 run.py                        |  6 ++++--
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/nd2reader/__init__.py b/nd2reader/__init__.py
index 7feb90e..2bfa070 100644
--- a/nd2reader/__init__.py
+++ b/nd2reader/__init__.py
@@ -1,6 +1,6 @@
 import logging
 from nd2reader.service import BaseNd2
-from nd2reader.model import Image
+from nd2reader.model import Image, ImageSet
 
 log = logging.getLogger("nd2reader")
 log.addHandler(logging.StreamHandler())
@@ -14,11 +14,12 @@ class Nd2(BaseNd2):
     def get_image(self, timepoint, fov, channel_name, z_level):
         image_set_number = self._calculate_image_set_number(timepoint, fov, z_level)
         timestamp, raw_image_data = self._reader.get_raw_image_data(image_set_number, self.channel_offset[channel_name])
-        return Image(timestamp, raw_image_data, self.height, self.width)
+        return Image(timestamp, raw_image_data, fov, channel_name, z_level, self.height, self.width)
 
     def __iter__(self):
         """
-        Just return every image.
+        Just return every image in order (might not be exactly the order that the images were physically taken, but it will
+        be within a few seconds). A better explanation is probably needed here.
 
         """
         for i in range(self._image_count):
@@ -27,4 +28,21 @@ class Nd2(BaseNd2):
                     for channel in self.channels:
                         image = self.get_image(i, fov, channel.name, z_level)
                         if image.is_valid:
-                            yield image
\ No newline at end of file
+                            yield image
+
+    def image_sets(self, field_of_view, timepoints=None, channels=None, z_levels=None):
+        """
+        Gets all the images for a given field of view and
+        """
+        timepoint_set = xrange(self._timepoint_count) if timepoints is None else timepoints
+        channel_set = [channel.name for channel in self.channels] if channels is None else channels
+        z_level_set = xrange(self._z_level_count) if z_levels is None else z_levels
+
+        for timepoint in timepoint_set:
+            image_set = ImageSet()
+            for channel_name in channel_set:
+                for z_level in z_level_set:
+                    image = self.get_image(timepoint, field_of_view, channel_name, z_level)
+                    if image.is_valid:
+                        image_set.add(image)
+            yield image_set
\ No newline at end of file
diff --git a/nd2reader/model/__init__.py b/nd2reader/model/__init__.py
index caf6f6b..56d0d90 100644
--- a/nd2reader/model/__init__.py
+++ b/nd2reader/model/__init__.py
@@ -1,5 +1,8 @@
 import numpy as np
 import skimage.io
+import logging
+
+log = logging.getLogger("nd2reader")
 
 
 class Channel(object):
@@ -43,23 +46,30 @@ class ImageSet(object):
 
 
 class Image(object):
-    def __init__(self, timestamp, raw_array, height, width):
+    def __init__(self, timestamp, raw_array, field_of_view, channel, z_level, height, width):
         self._timestamp = timestamp
         self._raw_data = raw_array
-        self._data = np.reshape(self._raw_data, (height, width))
+        self._field_of_view = field_of_view
+        self._channel = channel
+        self._z_level = z_level
+        self._height = height
+        self._width = width
+        self._data = None
 
     @property
     def timestamp(self):
         # TODO: Convert to datetime object
-        return self._timestamp
+        return self._timestamp / 1000.0
 
     @property
     def data(self):
+        if self._data is None:
+            self._data = np.reshape(self._raw_data, (self._height, self._width))
         return self._data
 
     @property
     def is_valid(self):
-        return np.any(self._data)
+        return np.any(self.data)
 
     def show(self):
         skimage.io.imshow(self.data)
diff --git a/nd2reader/service/__init__.py b/nd2reader/service/__init__.py
index ade86f9..bb7bd19 100644
--- a/nd2reader/service/__init__.py
+++ b/nd2reader/service/__init__.py
@@ -71,6 +71,10 @@ class BaseNd2(object):
         """
         return sum(self._metadata['ImageMetadata']['SLxExperiment']['ppNextLevelEx']['']['pItemValid'])
 
+    @property
+    def _channel_count(self):
+        return self._reader.channel_count
+
     @property
     def channel_offset(self):
         if self._channel_offset is None:
diff --git a/run.py b/run.py
index bd45aea..0c1f76b 100644
--- a/run.py
+++ b/run.py
@@ -4,9 +4,11 @@ import numpy as np
 from skimage import io
 
 
-for image in Nd2("/home/jim/Desktop/nd2hacking/YFP-dsRed-GFP-BF.nd2"):
+for image_set in Nd2("/home/jim/Desktop/nd2hacking/YFP-dsRed-GFP-BF.nd2").image_sets(0):
+    for image in image_set:
+        print(image.timestamp, image._channel, image._z_level)
+    print("")
 
 # n = Nd2("/home/jim/Desktop/nd2hacking/test-141111.nd2")
 
 
-    image.show()
\ No newline at end of file