diff --git a/pc.py b/pc.py
index 78211f1..ab7dbef 100644
--- a/pc.py
+++ b/pc.py
@@ -88,7 +88,8 @@ def read_images():
         if not os.path.exists(cell_fn):
             image = load_image(src_fn)
             image.resize((cell_w, cell_h), Image.LANCZOS).save(cell_fn)
-        if sum(data[i]) == 0:
+        #if sum(data[i]) == 0:
+        if np.all(data[i]==0):
             if not image:
                 image = load_image(src_fn)
             data[i] = np.array(image.resize((cx, cy), Image.LANCZOS).getdata())
@@ -106,6 +107,9 @@ def write_images(data):
     src_fns = [f'{src_dn}/{fn}' for fn in os.listdir(src_dn) if fn[0] != '.']
     src_fns = sorted(src_fns)
     n = len(src_fns)
+    if c > 1:
+        s = data.shape
+        data = data.reshape(s[0], s[1] * s[2])
     tree = cKDTree(data if kf == 1 else [data[i] for i in range(0, n, kf)])
     for i, src_fn in enumerate(src_fns):
         if i % ts != ti:
@@ -121,6 +125,10 @@ def write_images(data):
             for x in range(m):
                 crop = (x * cx, y * cy, (x + 1) * cx, (y + 1) * cy)
                 cell_data = src_image.crop(crop).getdata()
+                if c > 1:
+                    cell_data = np.array(cell_data)
+                    s = cell_data.shape
+                    cell_data = cell_data.reshape(s[0] * s[1])
                 cell_image = get_image(tree.query(cell_data)[1] * kf)
                 dst_image.paste(cell_image, (x * cell_w, y * cell_h))
         dst_image.save(dst_fn)