feat: add support for CloudVolume renumbered downsample

igneous/task_creation.py

@@ -235,7 +235,7 @@ def create_downsampling_tasks(
     sparse=False, bounds=None, chunk_size=None,
     encoding=None, delete_black_uploads=False, 
     background_color=0, dest_path=None, compress=None,
-    factor=None
+    factor=None, renumber=False
   ):
     """
     mip: Download this mip level, writes to mip levels greater than this one.
@@ -262,6 +262,9 @@ def create_downsampling_tasks(
       for new uploaded files.
     factor: (overrides axis) can manually specify what each downsampling round is
       supposed to do: e.g. (2,2,1), (2,2,2), etc
+    renumber: if True, use CloudVolume's renumbered download feature to download
+      a potentially smaller dtype so that larger images will fit in memory. Downloads
+      will be somewhat slower.
     """
     def ds_shape(mip, chunk_size=None, factor=None):
       if chunk_size:
@@ -306,6 +309,7 @@ def task(self, shape, offset):
           dest_path=dest_path,
           compress=compress,
           factor=factor,
+          renumber=bool(renumber),
         )
 
       def on_finish(self):
@@ -328,6 +332,7 @@ def on_finish(self):
             'dest_path': dest_path,
             'compress': compress,
             'factor': (tuple(factor) if factor else None),
+            'renumber': bool(renumber),
           },
           'by': OPERATOR_CONTACT,
           'date': strftime('%Y-%m-%d %H:%M %Z'),
@@ -900,7 +905,7 @@ def create_transfer_tasks(
     encoding=None, skip_downsamples=False,
     delete_black_uploads=False, background_color=0,
     agglomerate=False, timestamp=None, compress='gzip',
-    factor=None
+    factor=None, renumber=False
   ):
   """
   Transfer data from one data layer to another. It's possible
@@ -964,6 +969,7 @@ def task(self, shape, offset):
         timestamp=timestamp,
         compress=compress,
         factor=factor,
+        renumber=renumber,
       )
 
     def on_finish(self):
@@ -987,6 +993,7 @@ def on_finish(self):
           'timestamp': timestamp,
           'compress': compress,
           'factor': (tuple(factor) if factor else None),
+          'renumber': bool(renumber),
         },
         'by': OPERATOR_CONTACT,
         'date': strftime('%Y-%m-%d %H:%M %Z'),

igneous/tasks/tasks.py

@@ -39,7 +39,7 @@
 def downsample_and_upload(
     image, bounds, vol, ds_shape, 
     mip=0, axis='z', skip_first=False,
-    sparse=False, factor=None
+    sparse=False, factor=None, remap=None
   ):
     ds_shape = min2(vol.volume_size, ds_shape[:3])
     underlying_mip = (mip + 1) if (mip + 1) in vol.available_mips else mip
@@ -81,6 +81,8 @@ def downsample_and_upload(
       new_bounds //= factor3
       mipped = mips.pop(0)
       new_bounds.maxpt = new_bounds.minpt + Vec(*mipped.shape[:3])
+      if remap:
+        mipped = fastremap.remap(mipped, remap, in_place=True)
       vol[new_bounds] = mipped
 
 def cache(task, cloudpath):
@@ -699,7 +701,8 @@ def TransferTask(
   agglomerate=False,
   timestamp=None,
   compress='gzip',
-  factor=None
+  factor=None,
+  renumber=False
 ):
   shape = Vec(*shape)
   offset = Vec(*offset)
@@ -724,8 +727,14 @@ def TransferTask(
   dst_bounds = Bbox.clamp(dst_bounds, destcv.bounds)
   src_bounds = dst_bounds - translate
   image = srccv.download(
-    src_bounds, agglomerate=agglomerate, timestamp=timestamp
+    src_bounds, 
+    agglomerate=agglomerate, timestamp=timestamp, 
+    renumber=renumber
   )
+  remap = None
+  if renumber:
+    image, remap = image
+    remap = { v:k for k,v in remap.items() } # { label: 1, ... } -> { 1: label, ... }
 
   if skip_downsamples:
     destcv[dst_bounds] = image
@@ -735,15 +744,16 @@ def TransferTask(
       shape, mip=mip,
       skip_first=skip_first,
       sparse=sparse, axis=axis,
-      factor=factor
+      factor=factor, remap=remap
     )
 
 @queueable
 def DownsampleTask(
   layer_path, mip, shape, offset, 
   fill_missing=False, axis='z', sparse=False,
   delete_black_uploads=False, background_color=0,
-  dest_path=None, compress="gzip", factor=None
+  dest_path=None, compress="gzip", factor=None, 
+  renumber=False
 ):
   """
   Downsamples a cutout of the volume. By default it performs 
@@ -766,6 +776,7 @@ def DownsampleTask(
     axis=axis,
     compress=compress,
     factor=factor,
+    renumber=renumber,
   )
 
 class WatershedRemapTask(RegisteredTask):