Merge branch 'master' of github.com:seung-lab/chunkflow

Author: xiuliren

chunkflow/chunk/base.py

@@ -22,15 +22,15 @@
 # Offset = Tuple[int, int, int]
 from .validate import validate_by_template_matching
 
-def type_is_valid(type: str):
+def layer_type_is_valid(type: str):
     return type in set([None, 'image', 'segmentation', 'probability_map', 'affinity_map', 'unknown'])
 
 
 class Chunk(NDArrayOperatorsMixin):
     def __init__(self, array: np.ndarray, 
             voxel_offset: Cartesian = None, 
             voxel_size: Cartesian = None,
-            type: str = None):
+            layer_type: str = None):
         """chunk of a volume
     
         a chunk of big array with offset
@@ -49,7 +49,7 @@ def __init__(self, array: np.ndarray,
         """
         assert array.ndim >= 3 and array.ndim <= 4
         assert isinstance(array, np.ndarray) or isinstance(array, Chunk)
-        assert type_is_valid(type)
+        assert layer_type_is_valid(layer_type), f'layer type: {layer_type} is unsupported!'
 
         self.array = array
         if voxel_offset is None:
@@ -76,20 +76,20 @@ def __init__(self, array: np.ndarray,
             assert len(voxel_size) == 3
             assert np.alltrue([vs > 0 for vs in voxel_size])
 
-        if type is not None:
-            self.type = type 
+        if layer_type is not None:
+            self.layer_type = layer_type 
         else:
             # best guess
             if self.is_image:
-                self.type = 'image'
+                self.layer_type = 'image'
             elif self.is_segmentation:
-                self.type = 'segmentation'
+                self.layer_type = 'segmentation'
             elif self.is_probability_map:
-                self.type = 'probability_map'
+                self.layer_type = 'probability_map'
             elif self.is_affinity_map:
-                self.type = 'affinity_map'
+                self.layer_type = 'affinity_map'
             else:
-                self.type = 'unknown'
+                self.layer_type = 'unknown'
 
     # One might also consider adding the built-in list type to this
     # list, to support operations like np.add(array_like, list)
@@ -268,7 +268,7 @@ def from_h5(cls, file_name: str,
                 cutout_stop: tuple = None,
                 cutout_size: tuple = None,
                 dtype: str = None,
-                type: str = None):
+                layer_type: str = None):
 
         file_name = os.path.expanduser(file_name)
         if not os.path.exists(file_name):
@@ -310,17 +310,23 @@ def from_h5(cls, file_name: str,
                 else:
                     voxel_size = Cartesian(1, 1, 1)
 
-            if type is None:
-                if 'type' in f.attrs:
-                    type = f.attrs['type']
+            if layer_type is None:
+                if 'layer_type' in f.attrs:
+                    layer_type = f.attrs['layer_type']
                     # type = str(f['type'])
-                    breakpoint()
-                    assert type_is_valid(type)
+                    assert layer_type_is_valid(layer_type)
 
             if cutout_start is None:
                 cutout_start = voxel_offset
             if cutout_size is None:
                 cutout_size = dset.shape[-3:]
+                cutout_size = Cartesian.from_collection(cutout_size)
+            elif np.min(cutout_size) < 0:
+                cutout_size = [x for x in cutout_size]
+                for idx in range(-1, -4, -1):
+                    if cutout_size[idx]<0:
+                        cutout_size[idx] = dset.shape[idx]
+                cutout_size = Cartesian.from_collection(cutout_size)
             if cutout_stop is None:
                 cutout_stop = tuple(t+s for t, s in zip(cutout_start, cutout_size))
 
@@ -346,7 +352,7 @@ def from_h5(cls, file_name: str,
 
         logging.info(f'new chunk voxel offset: {cutout_start}')
 
-        return cls(arr, voxel_offset=cutout_start, voxel_size=voxel_size, type=type)
+        return cls(arr, voxel_offset=cutout_start, voxel_size=voxel_size, layer_type=layer_type)
 
     def to_h5(self, file_name: str, with_offset: bool=True, 
                 chunk_size: Union[Cartesian, tuple] = (8,8,8),
@@ -378,8 +384,8 @@ def to_h5(self, file_name: str, with_offset: bool=True,
                 voxel_size = self.voxel_size
             if voxel_size is not None:
                 f.create_dataset('/voxel_size', data=voxel_size)
-            if self.type is not None:
-                f.attrs['type'] = self.type
+            if self.layer_type is not None:
+                f.attrs['layer_type'] = self.layer_type
 
             if with_offset and self.voxel_offset is not None:
                 f.create_dataset('/voxel_offset', data=self.voxel_offset)
@@ -493,36 +499,15 @@ def is_affinity_map(self) -> bool:
     def is_probability_map(self) -> bool:
         return self.array.ndim == 4 and self.array.dtype == np.float32
 
-    # @property
-    # def type(self) -> str:
-    #     if self.data_type is None:
-    #         if self.is_image:
-    #             self.type = 'image'
-    #         elif self.is_segmentation:
-    #             self.type = 'segmentation'
-    #         elif self.is_probability_map:
-    #             self.type = 'probability_map'
-    #         elif self.is_affinity_map:
-    #             self.type = 'affinity_map'
-    #         else:
-    #             self.type = 'unknown'
-
-    #     return self.type
-
-    # @type.setter
-    # def type(self, value: str):
-    #     assert value in set([None, 'image', 'segmentation', 'probability_map', 'affinity_map', 'unknown'])
-    #     self.type = value
-    
     @property
     def properties(self) -> dict:
         props = dict()
         if self.voxel_offset is not None or self.voxel_offset != Cartesian(0, 0, 0):
             props['voxel_offset'] = self.voxel_offset
         if self.voxel_size is not None or self.voxel_size != Cartesian(1, 1, 1):
             props['voxel_size'] = self.voxel_size
-        if self.type is not None:
-            props['type'] = self.type
+        if self.layer_type is not None:
+            props['layer_type'] = self.layer_type
 
         return props 
 
@@ -533,8 +518,8 @@ def set_properties(self, properties: dict):
         if 'voxel_size' in properties:
             self.voxel_size = properties['voxel_size']
 
-        if 'type' in properties:
-            self.type = properties['type']
+        if 'layer_type' in properties:
+            self.layer_type = properties['layer_type']
 
     @properties.setter
     def properties(self, value: dict):
@@ -643,7 +628,7 @@ def channel_voting(self):
         return Chunk(out, 
             voxel_offset=self.voxel_offset, 
             voxel_size=self.voxel_size,
-            type='segmentation',
+            layer_type='segmentation',
         )
 
     def mask_using_last_channel(self, threshold: float = 0.3) -> np.ndarray:
@@ -742,7 +727,7 @@ def cutout(self, x: Union[tuple, BoundingBox, Bbox]):
         return Chunk(arr, 
             voxel_offset=voxel_offset, 
             voxel_size=self.voxel_size, 
-            type=self.type)
+            layer_type=self.layer_type)
 
     def save(self, patch):
         """

chunkflow/chunk/image/adjust_grey.py

@@ -199,7 +199,7 @@ def normalize_section_shang(image: np.ndarray, nominalmin: float,
     """
     assert nominalmin < nominalmax
     assert image.ndim == 3
-    voxel_offset = image.voxel_offset
+    # voxel_offset = image.voxel_offset
     originaltype = image.dtype
     arr = image.astype(np.float32)
 
@@ -228,8 +228,7 @@ def normalize_section_shang(image: np.ndarray, nominalmin: float,
     # cast to original data type if necessary
     #arr = np.round(arr)
     #arr = arr.astype(originaltype)
-
-    return Chunk(arr, voxel_offset=voxel_offset)
+    return arr
 
 
 def test1_grey_augment():

chunkflow/chunk/image/base.py

@@ -1,5 +1,6 @@
 __doc__ = """Image chunk class"""
 
+from tqdm import tqdm
 import numpy as np
 
 from chunkflow.chunk import Chunk
@@ -14,10 +15,119 @@ class Image(Chunk):
     def __init__(self, array: np.ndarray, **kwargs):
         super().__init__(array, **kwargs)
 
+    @classmethod
+    def from_chunk(cls, chk: Chunk):
+        return cls(chk.array, **chk.properties)
+
     def inference(self, inferencer: Inferencer):
         """run convolutional net inference for this image chunk"""
         return inferencer(self)
 
     def normalize_section_shang(self, nominalmin, nominalmax, clipvalues):
         return normalize_section_shang(self.array, nominalmin, nominalmax,
-                                       clipvalues)
+                                       clipvalues)
+
+    def _find_section_clamping_values(self, 
+            hist: np.ndarray, lower_clip_fraction: float, upper_clip_fraction: float):
+        """compute the clamping values for each section."""
+        # remove the np.copy from original code since we only need this once
+        filtered = hist
+
+        # remove pure black from frequency counts as
+        # it has no information in our images
+        filtered[0] = 0
+
+        cdf = np.zeros(shape=(len(filtered), ), dtype=np.uint64)
+        cdf[0] = filtered[0]
+        for i in range(1, len(filtered)):
+            cdf[i] = cdf[i - 1] + filtered[i]
+
+        total = cdf[-1]
+
+        if total == 0:
+            return (0, 0)
+
+        lower = 0
+        for i, val in enumerate(cdf):
+            if float(val) / float(total) > lower_clip_fraction:
+                break
+            lower = i
+
+        upper = 0
+        for i, val in enumerate(cdf):
+            if float(val) / float(total) > 1 - upper_clip_fraction:
+                break
+            upper = i
+        
+        return lower, upper
+    
+    def _hist_to_lookup_table(self, 
+            hist: np.ndarray, lower_clip_fraction: float, upper_clip_fraction: float,
+            minval: int = 1,
+            maxval: int = 255):
+        """histogram to lookup table
+
+        Args:
+            hist (np.ndarray): histogram
+
+        Returns:
+            np.ndarray: lookup table
+        """
+        lower, upper = self._find_section_clamping_values(
+            hist, lower_clip_fraction, upper_clip_fraction)
+             
+        if lower == upper:
+            #lookup_table = np.arange(0, 256, dtype=np.uint8)
+            # no need to perform any transform
+            return None
+        else:
+            # compute the lookup table
+            lookup_table = np.arange(0, 256, dtype=np.float32)
+            lookup_table = (lookup_table - float(lower)) * (
+                maxval / (float(upper) - float(lower)))
+            np.clip(lookup_table, minval, maxval, out=lookup_table)
+            lookup_table = np.round(lookup_table)
+            lookup_table = lookup_table.astype( np.uint8 ) 
+            return lookup_table
+
+
+    def normalize_contrast(self, 
+            lower_clip_fraction: float = 0.01, 
+            upper_clip_fraction: float = 0.01,
+            minval: int = 1,
+            maxval: int = 255,
+            per_section: bool = True
+            ):
+
+        def _normalize_array(array: np.ndarray, 
+                lower_clip_fraction: float, 
+                upper_clip_fraction: float,
+                minval: int = 1,
+                maxval: int = 255):
+            hist = np.bincount(array.flatten(), minlength=255)
+            lookup_table = self._hist_to_lookup_table(
+                hist, lower_clip_fraction, upper_clip_fraction, 
+                minval=minval, maxval=maxval)
+            if lookup_table is not None:
+                array = lookup_table[array]
+            return array
+        
+
+        if per_section:
+            for z in tqdm(range(self.bbox.start.z, self.bbox.stop.z)):
+                slices = (slice(z, z+1), *self.slices[-2:])
+                section = self.cutout(slices)
+                # section = Image.from_chunk(section)
+                section.array = _normalize_array(
+                    section.array,
+                    lower_clip_fraction, upper_clip_fraction,
+                    minval=minval, maxval=maxval
+                )
+                self.save(section)
+            else:
+                # chunk = Image.from_chunk(chunk)
+                self.array = _normalize_array(
+                    self.array,
+                    lower_clip_fraction, upper_clip_fraction,
+                    minval=minval, maxval=maxval 
+                )