From b294d732d1a127613befc9cc550a8a3f541851e4 Mon Sep 17 00:00:00 2001
From: Jingpeng Wu <jingpeng.wu@foxmail.com>
Date: Mon, 8 Apr 2024 15:44:01 +0800
Subject: [PATCH] fix AffinityMap destroy voxel size and voxel offset metadata
 bug

---
 chunkflow/chunk/affinity_map/base.py | 19 ++++++++++++++++---
 chunkflow/chunk/base.py              |  1 -
 chunkflow/flow/flow.py               |  2 +-
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/chunkflow/chunk/affinity_map/base.py b/chunkflow/chunk/affinity_map/base.py
index b21549aa..2693b50c 100644
--- a/chunkflow/chunk/affinity_map/base.py
+++ b/chunkflow/chunk/affinity_map/base.py
@@ -9,13 +9,26 @@ class AffinityMap(Chunk):
     """
     a chunk of affinity map. It has x,y,z three channels with single precision.
     """
-    def __init__(self, array, 
+    def __init__(self, array: np.ndarray, 
             voxel_offset: Cartesian=None, 
-            voxel_size: Cartesian=None ):
+            voxel_size: Cartesian=None,
+            layer_type: str = None):
+        assert isinstance(array, np.ndarray)
         assert array.ndim == 4
         assert np.issubdtype(array.dtype, np.float32)
         assert array.shape[0] == 3
-        super().__init__(array, voxel_offset=voxel_offset, voxel_size=voxel_size)
+        super().__init__(array, 
+            voxel_offset=voxel_offset, 
+            voxel_size=voxel_size, 
+            layer_type=layer_type)
+
+    @classmethod
+    def from_chunk(cls, chk: Chunk):
+        assert isinstance(chk, Chunk)
+        return cls(chk.array, 
+            voxel_offset = chk.voxel_offset,
+            voxel_size = chk.voxel_size,
+            layer_type = chk.layer_type)
 
     def quantize(self, mode: str='xy'):
         """transform affinity map to gray scale image
diff --git a/chunkflow/chunk/base.py b/chunkflow/chunk/base.py
index 7aeed965..c1b4ac1a 100755
--- a/chunkflow/chunk/base.py
+++ b/chunkflow/chunk/base.py
@@ -360,7 +360,6 @@ def from_h5(cls, file_name: str,
             arr = arr.astype('float64') 
 
         print(f'new chunk voxel offset: {cutout_start}')
-
         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, 
diff --git a/chunkflow/flow/flow.py b/chunkflow/flow/flow.py
index 30a869c4..40751175 100755
--- a/chunkflow/flow/flow.py
+++ b/chunkflow/flow/flow.py
@@ -2200,7 +2200,7 @@ def quantize(tasks, input_chunk_name: str, output_chunk_name: str, mode: str):
         if task is not None:
             chk = task[input_chunk_name]
             if chk.is_affinity_map:
-                chk = AffinityMap(chk)
+                chk = AffinityMap.from_chunk(chk)
                 quantized_image = chk.quantize(mode=mode)
             elif chk.is_probability_map:
                 quantized_image = (chk * 255.)