feat: continue to link up parts

Author: seankmartin

create_downsampled/chunking.py

@@ -105,6 +105,8 @@ def process(
     input_path,
     all_files,
     delete_input,
+    gcs_project,
+    num_channels,
 ):
     x_i, y_i, z_i = args
 
@@ -125,13 +127,22 @@ def process(
 
     # Use the new load_file function that handles download/caching
     print(f"Loading file for coordinates ({x_i}, {y_i}, {z_i})")
-    loaded_zarr_store = load_file(x_i, y_i)
+    loaded_zarr_store = load_file(
+        x_i,
+        y_i,
+        use_gcs_bucket,
+        input_path,
+        total_rows,
+        total_cols,
+        all_files,
+        gcs_project,
+    )
 
     if loaded_zarr_store is None:
         print(f"Warning: Could not load file for row {x_i}, col {y_i}. Skipping...")
         return
 
-    rawdata = load_data_from_zarr_store(loaded_zarr_store)
+    rawdata = load_data_from_zarr_store(loaded_zarr_store, num_channels)
 
     # Process all mip levels
     for mip_level in reversed(range(mip_cutoff, num_mips)):

create_downsampled/gcs_local_io.py

@@ -225,6 +225,7 @@ def check_and_upload_completed_chunks(
     use_gcs_output,
     gcs_project,
     gcs_output_bucket_name,
+    gcs_output_path,
     num_upload_workers,
     delete_output,
     already_uploaded_path,
@@ -283,6 +284,10 @@ def check_and_upload_completed_chunks(
                 gcs_output_bucket_name,
                 files_to_upload_this_batch,
                 source_directory=output_path,
+                gcs_output_path=gcs_output_path,
+                gcs_project=gcs_project,
+                uploaded_files=uploaded_files,
+                failed_files=failed_files,
                 workers=num_upload_workers,
             )
             uploaded_count += len(files_to_upload_this_batch)

create_downsampled/main.py

@@ -4,7 +4,12 @@
 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 gcs_local_io import load_file, check_and_upload_completed_chunks, check_any_remaining_chunks
+from gcs_local_io import (
+    get_uploaded_files,
+    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
 
@@ -89,6 +94,9 @@ def main():
     # Process each well into chunks
     iter_coords = list(get_grid_coords(num_chunks_per_dim))
 
+    # Find which files were already done and keep track of them
+    uploaded_files = get_uploaded_files(output_path)
+
     processed_chunks = []
     failed_chunks = []
     total_uploads = 0
@@ -110,18 +118,36 @@ def main():
             input_path=input_path,
             all_files=all_files,
             delete_input=delete_input,
+            gcs_project=gcs_project,
+            num_channels=num_channels,
         )
         start, end = bounds
         processed_chunks.append((start, end))
-        total_uploads += check_and_upload_completed_chunks()
+        total_uploads += check_and_upload_completed_chunks(
+            num_mips=num_mips,
+            output_path=output_path,
+            volume_size=volume_size,
+            processed_chunks_bounds=processed_chunks,
+            use_gcs_output=use_gcs_output,
+            gcs_project=gcs_project,
+            gcs_output_bucket_name=gcs_output_bucket_name,
+            gcs_output_path=gcs_output_path,
+            num_upload_workers=num_upload_workers,
+            delete_output=delete_output,
+            already_uploaded_path=output_path / "uploaded_to_gcs_chunks.txt",
+            uploaded_files=uploaded_files,
+            failed_files=failed_chunks,
+        )
         print(f"Total chunks uploaded so far: {total_uploads}")
 
     if failed_chunks:
         print(f"Failed to process {len(failed_chunks)} chunks:")
         for chunk in failed_chunks:
             print(f"  {chunk}")
 
-    remaining_files = check_any_remaining_chunks()
+    remaining_files = check_any_remaining_chunks(
+        num_mips=num_mips, output_path=output_path, uploaded_files=uploaded_files
+    )
     if remaining_files:
         print(f"Remaining chunks: {remaining_files}")