prepare datasets
2 changed files+113−25
prepare_datasets/aind/list_s3_directories.pyadded+77−0View file
@@ -0,0 +1,77 @@
1+#!/usr/bin/env python3
2+"""
3+Script to list directories in s3://aind-benchmark-data/ephys-compression
4+Lists directories up to depth 2
5+"""
6+import boto3
7+from botocore import UNSIGNED
8+from botocore.config import Config
9+
10+
11+def list_s3_directories(bucket_name: str, prefix: str = '', depth: int = 1):
12+ """
13+ List directories (common prefixes) in an S3 bucket recursively up to a certain depth
14+
15+ Args:
16+ bucket_name: Name of the S3 bucket
17+ prefix: Prefix path to list directories under (default: root)
18+ depth: How deep to recurse (1 = only immediate subdirs)
19+ """
20+ # S3 configuration for public bucket (no credentials needed)
21+ s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
22+
23+ # Ensure prefix ends with / if not empty
24+ if prefix and not prefix.endswith('/'):
25+ prefix += '/'
26+
27+ # List objects with delimiter to get "directories"
28+ paginator = s3.get_paginator('list_objects_v2')
29+ pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter='/')
30+
31+ directories = []
32+
33+ for page in pages:
34+ # CommonPrefixes contains the "directories"
35+ if 'CommonPrefixes' in page:
36+ for prefix_obj in page['CommonPrefixes']:
37+ dir_path = prefix_obj['Prefix']
38+ # Remove the base prefix to get relative path
39+ relative_dir = dir_path[len(prefix):] if prefix else dir_path
40+ directories.append(relative_dir)
41+
42+ # Sort directories
43+ directories = sorted(directories)
44+
45+ # Print current level directories and their subdirectories
46+ for directory in directories:
47+ print(f" {directory}")
48+
49+ # Recurse if depth > 1
50+ if depth > 1:
51+ # List subdirectories
52+ full_prefix = prefix + directory
53+ subpages = paginator.paginate(Bucket=bucket_name, Prefix=full_prefix, Delimiter='/')
54+
55+ subdirs = []
56+ for page in subpages:
57+ if 'CommonPrefixes' in page:
58+ for prefix_obj in page['CommonPrefixes']:
59+ dir_path = prefix_obj['Prefix']
60+ relative_dir = dir_path[len(full_prefix):] if full_prefix else dir_path
61+ subdirs.append(relative_dir)
62+
63+ # Print subdirectories with indentation
64+ for subdir in sorted(subdirs):
65+ print(f" {subdir}")
66+
67+ return directories
68+
69+
70+if __name__ == '__main__':
71+ bucket_name = 'aind-benchmark-data'
72+ prefix = 'ephys-compression'
73+
74+ print(f"Listing directories in s3://{bucket_name}/{prefix} (depth 2)")
75+ print()
76+
77+ list_s3_directories(bucket_name, prefix, depth=2)
prepare_datasets/aind/prepare_aind_compression_np2_probeB.pymodified+36−25View file
@@ -3,32 +3,43 @@ import spikeinterface as si
33 import numpy as np
44 from s3_utils import download_s3_folder
55
6-s3_folder_name = "s3://aind-benchmark-data/ephys-compression/aind-np2/612962_2022-04-13_19-18-04_ProbeB"
7-local_folder_name = "612962_2022-04-13_19-18-04_ProbeB.si"
6+s3_base_url = "s3://aind-benchmark-data/ephys-compression/"
87
9-if not os.path.exists(local_folder_name):
10- download_s3_folder(s3_folder_name, local_folder_name)
8+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")
11+]
1112
12-recording = si.load(
13- local_folder_name
14-)
13+for folder_name, name0 in folder_names:
14+ s3_folder_name = f"{s3_base_url}/{folder_name}"
15+ local_folder_name = f"{folder_name}.si"
1516
16-channel_ids = [
17- 'CH101',
18- 'CH102',
19- 'CH103',
20- 'CH104',
21- 'CH105',
22- 'CH106',
23- 'CH107',
24- 'CH108',
25- 'CH109',
26- 'CH110'
27-]
17+ if not os.path.exists(local_folder_name):
18+ # make parent directories if needed
19+ os.makedirs(os.path.dirname(local_folder_name), exist_ok=True)
20+ print(f'Downloading {s3_folder_name} to {local_folder_name}...')
21+ download_s3_folder(s3_folder_name, local_folder_name)
22+
23+ recording = si.load(
24+ local_folder_name
25+ )
26+
27+ channel_ids = [
28+ 'CH101',
29+ 'CH102',
30+ 'CH103',
31+ 'CH104',
32+ 'CH105',
33+ 'CH106',
34+ 'CH107',
35+ 'CH108',
36+ 'CH109',
37+ 'CH110'
38+ ]
2839
29-fname = f'aind_compression_np2_probeB_ch101-110.raw.npy'
30-if not os.path.exists(fname):
31- print(f'Writing {fname}...')
32- X = recording.get_traces(channel_ids=channel_ids, start_frame=30000, end_frame=30000 + 30000 * 10)
33- print(f'X.shape = {X.shape}')
34- np.save(fname, X)
40+ fname = f'{name0}-ch101-110.raw.npy'
41+ if not os.path.exists(fname):
42+ print(f'Writing {fname}...')
43+ X = recording.get_traces(channel_ids=channel_ids, start_frame=30000, end_frame=30000 + 30000 * 10)
44+ print(f'X.shape = {X.shape}')
45+ np.save(fname, X)