prepare datasets
3 changed files+22−16
prepare_datasets/.gitignoremodified+3−1View file
@@ -1,2 +1,4 @@
11 *.si
2-*.npy
2+*.npy
3+
4+__pycache__
prepare_datasets/aind/prepare_aind_compression_np2_probeB.py →prepare_datasets/aind/prepare_aind_compression_datasets.pyrenamed+16−14View file
@@ -3,14 +3,14 @@ import spikeinterface as si
33 import numpy as np
44 from s3_utils import download_s3_folder
55
6-s3_base_url = "s3://aind-benchmark-data/ephys-compression/"
6+s3_base_url = "s3://aind-benchmark-data/ephys-compression"
77
88 folder_names = [
9- ("aind-np2/612962_2022-04-13_19-18-04_ProbeB", "aind-np2-probeB"),
10- ("aind-np1/625749_2022-08-03_15-15-06_ProbeA", "aind-np1-probeA")
9+ # ("aind-np2/612962_2022-04-13_19-18-04_ProbeB", "aind-np2-probeB", "CH"),
10+ ("aind-np1/625749_2022-08-03_15-15-06_ProbeA", "aind-np1-probeA", "AP")
1111 ]
1212
13-for folder_name, name0 in folder_names:
13+for folder_name, name0, channel_prefix in folder_names:
1414 s3_folder_name = f"{s3_base_url}/{folder_name}"
1515 local_folder_name = f"{folder_name}.si"
1616
@@ -18,23 +18,25 @@ for folder_name, name0 in folder_names:
1818 # make parent directories if needed
1919 os.makedirs(os.path.dirname(local_folder_name), exist_ok=True)
2020 print(f'Downloading {s3_folder_name} to {local_folder_name}...')
21+ # IMPORTANT NOTE: we may interrupt this download early because we really only need the first part.
2122 download_s3_folder(s3_folder_name, local_folder_name)
2223
24+ # For now this only works with spikeinterface 0.102
2325 recording = si.load(
2426 local_folder_name
2527 )
2628
2729 channel_ids = [
28- 'CH101',
29- 'CH102',
30- 'CH103',
31- 'CH104',
32- 'CH105',
33- 'CH106',
34- 'CH107',
35- 'CH108',
36- 'CH109',
37- 'CH110'
30+ f'{channel_prefix}101',
31+ f'{channel_prefix}102',
32+ f'{channel_prefix}103',
33+ f'{channel_prefix}104',
34+ f'{channel_prefix}105',
35+ f'{channel_prefix}106',
36+ f'{channel_prefix}107',
37+ f'{channel_prefix}108',
38+ f'{channel_prefix}109',
39+ f'{channel_prefix}110'
3840 ]
3941
4042 fname = f'{name0}-ch101-110.raw.npy'
prepare_datasets/aind/s3_utils.pymodified+3−1View file
@@ -115,9 +115,11 @@ def download_s3_folder(s3_url: str, local_dir: str, skip_confirmation: bool = Fa
115115 Path(local_file_dir).mkdir(parents=True, exist_ok=True)
116116
117117 # Download the file with progress callback
118+ # Using download_fileobj to write directly to file for partial download support
118119 print(f"[{idx}/{len(files_to_download)}] {relative_path} ({file_size:,} bytes)")
119120 progress = ProgressCallback(relative_path, file_size)
120- s3.download_file(bucket_name, s3_key, local_file, Callback=progress)
121+ with open(local_file, 'wb') as f:
122+ s3.download_fileobj(bucket_name, s3_key, f, Callback=progress)
121123 print() # New line after progress
122124
123125 downloaded_bytes += file_size