feat: continue with refactored processing

Author: seankmartin

create_downsampled/.env.example

@@ -0,0 +1,32 @@
+# Example environment configuration file
+# Copy this file to .env and modify the values as needed
+
+# Data source configuration
+USE_GCS_BUCKET=true
+GCS_BUCKET_NAME=YOUR_GCS_BUCKET_NAME
+GCS_PREFIX=YOUR/GCS/PREFIX
+GCS_FILE_EXTENSION=.zarr
+GCS_PROJECT=YOUR_GCP_PROJECT_ID
+GOOGLE_CLOUD_PROJECT=YOUR_GCP_PROJECT_ID
+GCS_FILES_LOCAL_LIST=/path/to/bucket-list.txt
+
+# Output GCS bucket configuration (for uploading processed results)
+USE_GCS_OUTPUT=true
+GCS_OUTPUT_BUCKET_NAME=YOUR_OUTPUT_GCS_BUCKET_NAME
+GCS_OUTPUT_PREFIX=YOUR/OUTPUT/PREFIX
+OVERWRITE_GCS=false
+NUM_UPLOAD_WORKERS=8
+
+# Local paths (used when USE_GCS_BUCKET=false)
+INPUT_PATH=/path/to/input
+OUTPUT_PATH=/path/to/output
+DELETE_INPUT=false
+DELETE_OUTPUT=false
+
+# Processing settings
+OVERWRITE=false
+NUM_MIPS=5
+MIP_CUTOFF=3
+CHANNEL_LIMIT=4
+ALLOW_NON_ALIGNED_WRITE=true
+MANUAL_CHUNK_SIZE=64,64,32

create_downsampled/chunk_utils.py

@@ -0,0 +1,49 @@
+def is_chunk_fully_covered(chunk_bounds, processed_chunks_bounds):
+    """
+    Check if a chunk is fully covered by processed bounds.
+
+    Args:
+        chunk_bounds: [start_coord, end_coord] where each coord is [x, y, z]
+        processed_chunks_bounds: List of tuples (start, end) where start and end are [x, y, z]
+
+    Returns:
+        bool: True if all 8 corners of the chunk are covered by processed bounds
+    """
+    if not processed_chunks_bounds:
+        return False
+
+    start_coord, end_coord = chunk_bounds
+    x0, y0, z0 = start_coord
+    x1, y1, z1 = end_coord
+
+    # Generate all 8 corners of the chunk
+    corners = [
+        [x0, y0, z0],  # min corner
+        [x1, y0, z0],
+        [x0, y1, z0],
+        [x0, y0, z1],
+        [x1, y1, z0],
+        [x1, y0, z1],
+        [x0, y1, z1],
+        [x1, y1, z1],  # max corner
+    ]
+
+    # Check if each corner is covered by at least one processed bound
+    for corner in corners:
+        corner_covered = False
+        for start, end in processed_chunks_bounds:
+            # Check if corner is inside this processed bound
+            if (
+                start[0] <= corner[0] < end[0]
+                and start[1] <= corner[1] < end[1]
+                and start[2] <= corner[2] < end[2]
+            ):
+                corner_covered = True
+                break
+
+        # If any corner is not covered, the chunk is not fully covered
+        if not corner_covered:
+            return False
+
+    # All corners are covered
+    return True

create_downsampled/chunking.py

@@ -4,8 +4,7 @@
 import numpy as np
 from neuroglancer.downsample import downsample_with_averaging
 
-from io import load_data_from_zarr_store, delete_cached_zarr_file, load_file
-
+from gcs_local_io import load_data_from_zarr_store, delete_cached_zarr_file, load_file
 
 
 def compute_optimal_chunk_size(single_file_shape, num_mips, max_chunk_size=None):
@@ -100,6 +99,12 @@ def process(
     allow_non_aligned_write,
     overwrite_output,
     progress_dir,
+    total_rows,
+    total_cols,
+    use_gcs_bucket,
+    input_path,
+    all_files,
+    delete_input,
 ):
     x_i, y_i, z_i = args
 
@@ -202,58 +207,16 @@ def process(
     touch(f_name)
 
     # Clean up cached file to save disk space
-    delete_cached_zarr_file(x_i, y_i)
+    delete_cached_zarr_file(
+        x_i,
+        y_i,
+        total_rows,
+        total_cols,
+        use_gcs_bucket,
+        delete_input,
+        input_path,
+        all_files,
+    )
 
     # Return the bounds of the processed chunk
     return (start, end)
-
-
-def is_chunk_fully_covered(chunk_bounds, processed_chunks_bounds):
-    """
-    Check if a chunk is fully covered by processed bounds.
-
-    Args:
-        chunk_bounds: [start_coord, end_coord] where each coord is [x, y, z]
-        processed_chunks_bounds: List of tuples (start, end) where start and end are [x, y, z]
-
-    Returns:
-        bool: True if all 8 corners of the chunk are covered by processed bounds
-    """
-    if not processed_chunks_bounds:
-        return False
-
-    start_coord, end_coord = chunk_bounds
-    x0, y0, z0 = start_coord
-    x1, y1, z1 = end_coord
-
-    # Generate all 8 corners of the chunk
-    corners = [
-        [x0, y0, z0],  # min corner
-        [x1, y0, z0],
-        [x0, y1, z0],
-        [x0, y0, z1],
-        [x1, y1, z0],
-        [x1, y0, z1],
-        [x0, y1, z1],
-        [x1, y1, z1],  # max corner
-    ]
-
-    # Check if each corner is covered by at least one processed bound
-    for corner in corners:
-        corner_covered = False
-        for start, end in processed_chunks_bounds:
-            # Check if corner is inside this processed bound
-            if (
-                start[0] <= corner[0] < end[0]
-                and start[1] <= corner[1] < end[1]
-                and start[2] <= corner[2] < end[2]
-            ):
-                corner_covered = True
-                break
-
-        # If any corner is not covered, the chunk is not fully covered
-        if not corner_covered:
-            return False
-
-    # All corners are covered
-    return True

create_downsampled/io.py renamed to create_downsampled/gcs_local_io.py

@@ -6,12 +6,12 @@
 import zarr
 import numpy as np
 
-from chunking import is_chunk_fully_covered
+from chunk_utils import is_chunk_fully_covered
 from wells import extract_row_col_from_filename
 from gcs import gcloud_download_dir, upload_many_blobs_with_transfer_manager
 
 
-def get_local_cache_path(row, col, use_gcs_bucket, input_path):
+def get_local_cache_path(remote_file, use_gcs_bucket, input_path):
     """
     Get the local cache path where a file for the given row/col should be stored.
 
@@ -23,11 +23,6 @@ def get_local_cache_path(row, col, use_gcs_bucket, input_path):
         Path: Local cache path for the file
     """
 
-    # Get the remote file path/name for this row/col
-    remote_file = get_remote_file_path(row, col)
-    if remote_file is None:
-        return None
-
     # Create local filename based on remote file
     if use_gcs_bucket:
         cache_dir = input_path
@@ -83,7 +78,7 @@ def download_zarr_file(
         print(f"No file found for row {row}, col {col}")
         return None
 
-    local_path = get_local_cache_path(row, col, use_gcs_bucket, input_path)
+    local_path = get_local_cache_path(remote_file, use_gcs_bucket, input_path)
     if local_path is None:
         return None
 
@@ -137,7 +132,16 @@ def load_file(
         return None
 
 
-def delete_cached_zarr_file(row, col, use_gcs_bucket, delete_input, input_path):
+def delete_cached_zarr_file(
+    row,
+    col,
+    total_rows,
+    total_cols,
+    use_gcs_bucket,
+    delete_input,
+    input_path,
+    all_files,
+):
     """
     Delete the locally cached file for a specific row and column to save disk space.
 
@@ -150,7 +154,8 @@ def delete_cached_zarr_file(row, col, use_gcs_bucket, delete_input, input_path):
     """
     if not use_gcs_bucket or not delete_input:
         return True
-    local_path = get_local_cache_path(row, col, use_gcs_bucket, input_path)
+    remote_path = get_remote_file_path(row, col, total_rows, total_cols, all_files)
+    local_path = get_local_cache_path(remote_path, use_gcs_bucket, input_path)
     if local_path is None:
         return True
 

create_downsampled/main.py

@@ -4,7 +4,7 @@
 from load_config import load_env_config
 from gcs import list_gcs_files, sync_info_to_gcs_output
 from wells import compute_grid_dimensions, get_grid_coords
-from io import load_file, check_and_upload_completed_chunks, check_any_remaining_chunks
+from gcs_local_io import load_file, check_and_upload_completed_chunks, check_any_remaining_chunks
 from chunking import compute_volume_and_chunk_size, process
 from volume import create_cloudvolume_info
 
@@ -87,28 +87,29 @@ def main():
     )
 
     # Process each well into chunks
-    iter_coords = get_grid_coords(num_chunks_per_dim)
+    iter_coords = list(get_grid_coords(num_chunks_per_dim))
 
     processed_chunks = []
     failed_chunks = []
     total_uploads = 0
     for coord in iter_coords[:MAX_ITERS]:
         bounds = process(
-            coord,
-            single_file_shape,
-            volume_shape,
-            vols,
-            chunk_size,
-            num_mips,
-            mip_cutoff,
-            allow_non_aligned_write,
-            overwrite_output,
-            use_gcs_bucket,
-            input_path,
-            total_rows,
-            total_cols,
-            all_files,
-            gcs_project,
+            args=coord,
+            single_file_shape=single_file_shape,
+            volume_shape=volume_shape,
+            vols=vols,
+            chunk_size=chunk_size,
+            num_mips=num_mips,
+            mip_cutoff=mip_cutoff,
+            allow_non_aligned_write=allow_non_aligned_write,
+            overwrite_output=overwrite_output,
+            progress_dir=output_path / "progress",
+            total_rows=total_rows,
+            total_cols=total_cols,
+            use_gcs_bucket=use_gcs_bucket,
+            input_path=input_path,
+            all_files=all_files,
+            delete_input=delete_input,
         )
         start, end = bounds
         processed_chunks.append((start, end))
@@ -133,3 +134,7 @@ def main():
         gcs_project,
         gcs_output_bucket_name,
     )
+
+
+if __name__ == "__main__":
+    main()