import pickle
import shutil
from collections.abc import Iterator
from contextlib import contextmanager
from pathlib import Path
from typing import Any, TypeAlias
import gdown
import numpy as np
import polars as pl
from neuralib.io import NEUROLIB_DATASET_DIRECTORY
from neuralib.typing import PathLike
from PIL import Image
DeepLabCutDataFrame: TypeAlias = Any
Suite2PResult: TypeAlias = Any
SBXInfo: TypeAlias = Any
FaceMapResult: TypeAlias = Any
RasterMapResult: TypeAlias = Any
CziScanner: TypeAlias = Any
TiffScanner: TypeAlias = Any
__all__ = [
'google_drive_file',
'google_drive_folder',
'clean_all_cache_dataset',
#
'load_example_rois',
'load_example_rois_image',
'load_example_rois_dir',
'load_example_dorsal_cortex',
#
'load_ephys_meta',
'load_ephys_data',
'load_npx2_reconstruction',
#
'load_example_scanbox',
'load_example_suite2p_result',
'load_example_rastermap_2p_result',
'load_example_rastermap_2p_cache',
'load_example_retinotopic_data',
#
'load_example_dlc_h5',
'load_example_dlc_csv',
'load_example_facemap_pupil',
'load_example_facemap_keypoints',
#
'load_example_lsm',
'load_example_czi'
]
[docs]
@contextmanager
def google_drive_file(file_id: str,
*,
quiet: bool = False,
output_dir: PathLike | None = None,
rename_file: str | None = None,
cached: bool = False,
invalid_cache: bool = False) -> Iterator[Path]:
"""
Download file from Google Drive. If not ``cached``, then delete afterward.
:param file_id: Google Drive file ID used to identify the file to be downloaded.
:param quiet: Boolean flag to suppress output from the download process.
:param output_dir: Directory path where the downloaded file will be saved.
:param rename_file: Optional string to rename the downloaded file.
:param cached: Boolean flag to retain the downloaded file after usage.
:param invalid_cache: Boolean flag to force re-download even if the file exists in the cache.
:return: A context manager yielding the path to the downloaded file.
"""
if rename_file is not None:
file = rename_file
else:
file = file_id
if output_dir is None:
output_dir = NEUROLIB_DATASET_DIRECTORY
else:
output_dir = Path(output_dir)
output_dir.mkdir(exist_ok=True, parents=True)
output_file = output_dir / file
try:
if output_file.exists() and not invalid_cache:
yield output_file
else:
url = f"https://drive.google.com/uc?id={file_id}"
gdown.download(url, str(output_file), quiet=quiet) # pyright: ignore[reportPrivateImportUsage]
yield output_file
finally:
if not cached:
output_file.unlink(missing_ok=True)
[docs]
@contextmanager
def google_drive_folder(folder_id: str,
*,
quiet: bool = False,
output_dir: Path | None = None,
rename_folder: str | None = None,
cached: bool = False,
invalid_cache: bool = False) -> Iterator[Path]:
"""
Download a entire folder from Google Drive. If not ``cached``, then delete afterward.
:param folder_id: Google Drive folder ID to download.
:param quiet: If True, suppress the gdown output.
:param output_dir: Directory where the folder should be downloaded. If None, a temporary directory is used.
:param rename_folder: Optional new name for the downloaded folder.
:param cached: If True, keep the downloaded folder for future use.
:param invalid_cache: If True, force re-download even if cached data exists.
:return: A context manager yielding the path to the downloaded folder.
"""
if rename_folder is not None:
folder_name = rename_folder
else:
folder_name = folder_id
if output_dir is None:
output_dir = NEUROLIB_DATASET_DIRECTORY / folder_name
try:
if output_dir.exists() and any(output_dir.iterdir()) and not invalid_cache:
yield output_dir
else:
output_dir.mkdir(exist_ok=True, parents=True)
gdown.download_folder(id=folder_id, output=str(output_dir), quiet=quiet) # pyright: ignore[reportPrivateImportUsage]
yield output_dir
finally:
if not cached:
shutil.rmtree(output_dir, ignore_errors=True)
[docs]
def clean_all_cache_dataset():
"""removes the cache dataset directory"""
shutil.rmtree(NEUROLIB_DATASET_DIRECTORY)
print(f'[REMOVE]: {NEUROLIB_DATASET_DIRECTORY}')
# ========== #
# Atlas Data #
# ========== #
[docs]
def load_example_rois(**kwargs) -> pl.DataFrame:
"""
Load example ROIs dataframe
:param kwargs: Additional keyword arguments pass to ``google_drive_file`` to customize the loading behavior.
:return: A Polars DataFrame containing the example ROIs data
"""
with google_drive_file('1dKpZt6eF4szvl7svWRdBQkOfTVLQi4Xg', **kwargs) as file:
return pl.read_csv(file)
[docs]
def load_example_rois_image(**kwargs) -> np.ndarray:
"""
Load example ROIs image array
:param kwargs: Additional arguments to be passed to the `google_drive_file` context manager.
:return: An example imaging array with labeled ROIs
"""
with google_drive_file('1-ZFC7Fd6IgwbY6X8oetvpGjZaBBiwg0t', **kwargs) as file:
return np.array(Image.open(file))
[docs]
def load_example_rois_dir(**kwargs) -> Path:
"""Load a directory containing multiple example ROIs image array"""
with google_drive_folder('1tj36-lzOArjwMsZlhbyy_B19ffl50J98', **kwargs) as d:
return d
[docs]
def load_example_dorsal_cortex(color: bool = False, **kwargs) -> np.ndarray:
"""Load dorsal cortex image array from the source svg
.. seealso::
:meth:`~neuralib.atlas.data.get_dorsal_cortex`
:param color:
"""
if color:
file_id = '1Cujx3GGFZxEq0-isRlA_Ac-Puy6ml8IA'
else:
file_id = '1OEPpIIl8SszDJXO_1ADka-5pB556WNxc'
with google_drive_file(file_id, **kwargs) as file:
return np.array(Image.open(file))
# ========== #
# Ephys Data #
# ========== #
[docs]
def load_ephys_data(**kwargs):
with google_drive_file('1U0xAchQagyXRT72M68fRQ4JsRQeW9q5d', **kwargs) as _:
pass
[docs]
def load_npx2_reconstruction(**kwargs) -> pl.DataFrame:
"""
Example of NeuroPixel2 reconstruction data
:param kwargs: Additional keyword arguments pass to ``google_drive_file`` to customize the loading behavior.
:return: A Polars DataFrame containing the example DiI labelled traces ROIs
"""
with google_drive_file('1fRvMNHhGgh5KP3CgGm6CMFth1qIAmwfh', **kwargs) as file:
return pl.read_csv(file)
# ============ #
# Imaging Data #
# ============ #
[docs]
def load_example_scanbox(**kwargs) -> 'SBXInfo':
"""
Load example ScanBox data
:param kwargs: Additional keyword arguments pass to ``google_drive_file`` to customize the loading behavior.
:return: An instance of ``SBXInfo`` loaded from the specified Google Drive file.
"""
from neuralib.scanbox import SBXInfo # pyright: ignore[reportMissingImports]
with google_drive_file('1Gcz_xRVWQJ9QMxq3vzZS8VruSbNiuh_s', **kwargs) as file:
return SBXInfo.load(file)
[docs]
def load_example_suite2p_result(**kwargs) -> 'Suite2PResult':
"""
Load example Suite2P data
:param kwargs: Additional keyword arguments pass to ``google_drive_folder`` to customize the loading behavior.
:return: An instance of ``Suite2PResult`` loaded with data from the specified Google Drive folder.
"""
from neuralib.suite2p import read_suite2p # pyright: ignore[reportMissingImports]
with google_drive_folder('1iVImr_rIywWhCiBDYhcphcSODaWJrhy7', **kwargs) as suite2p_dir:
return read_suite2p(suite2p_dir)
[docs]
def load_example_rastermap_2p_result(**kwargs) -> 'RasterMapResult':
"""Load example rastermap 2P data"""
from neuralib.rastermap import read_rastermap # pyright: ignore[reportMissingImports]
with google_drive_file('1KSic4sXyF3hTgQbGijMpa3D3TGJUU097', **kwargs) as file:
return read_rastermap(file)
[docs]
def load_example_rastermap_2p_cache(**kwargs) -> dict[str, Any]:
"""
Load example rastermap 2P cache pickle file
:param kwargs: Additional arguments to be passed to the `google_drive_file` context manager.
:return: A dictionary containing the 2-photon rastermap data cache
"""
with google_drive_file('1SuzUhkbcdBY71dCnxjnT4pmQOsiOULBC', **kwargs) as file:
with file.open('rb') as f:
return pickle.load(f)
[docs]
def load_example_retinotopic_data(**kwargs) -> Path:
"""Load example retinotopic tiff file path"""
with google_drive_file('1J8iqP_EBaknNJehRUw3nwp7lEw2_UFXz',
quiet=True,
cached=True,
rename_file='retinotopic.tiff', **kwargs) as file:
return Path(file)
# ========== #
# Behavioral #
# ========== #
[docs]
def load_example_dlc_h5(**kwargs) -> 'DeepLabCutDataFrame':
"""Load example Deeplabcut h5"""
with google_drive_file('1JNhx6Dpe8beP8vnh0yF3o3vY2DfUM-8A', rename_file='test_dlc.h5', **kwargs) as h5:
from neuralib.deeplabcut.core import read_dlc # pyright: ignore[reportMissingImports]
with google_drive_file('1juICYcrXa7Vk-fQSBBSg2QcP9DGyHO2E', rename_file='test_dlc.pickle', **kwargs) as meta:
return read_dlc(h5, meta)
[docs]
def load_example_dlc_csv(**kwargs) -> 'DeepLabCutDataFrame':
"""Load example Deeplabcut csv"""
from neuralib.deeplabcut.core import read_dlc # pyright: ignore[reportMissingImports]
with google_drive_file('1R2Ze5xjWlavcKvu45JOH3_QkOD4SSkVN', rename_file='test_dlc.csv', **kwargs) as csv:
with google_drive_file('1juICYcrXa7Vk-fQSBBSg2QcP9DGyHO2E', rename_file='test_dlc.pickle', **kwargs) as meta:
return read_dlc(csv, meta)
[docs]
def load_example_facemap_pupil(**kwargs) -> 'FaceMapResult':
"""Load example facemap pupil data"""
from neuralib.facemap import read_facemap # pyright: ignore[reportMissingImports]
with google_drive_folder('1cacH5DWLmYqh_7PLwqEasmER_TfKgZ1b', **kwargs) as pupil_dir:
return read_facemap(pupil_dir)
[docs]
def load_example_facemap_keypoints(**kwargs) -> 'FaceMapResult':
"""Load example facemap keypoint data"""
from neuralib.facemap import read_facemap # pyright: ignore[reportMissingImports]
with google_drive_folder('1FWz70HE_hQuhE6K9hoO_y1OgeG11NsGM', **kwargs) as pupil_dir:
return read_facemap(pupil_dir)
# ============= #
# Confocal Scan #
# ============= #
[docs]
def load_example_lsm(**kwargs) -> 'TiffScanner':
"""load example lsm file"""
from neuralib.scan.lsm import lsm_file # pyright: ignore[reportMissingImports]
with google_drive_file('1beq6PCY8XmZjyWiOk2-KkcXLrmbLpWmS', rename_file='test.lsm', **kwargs) as file:
with lsm_file(file) as lsm:
return lsm
[docs]
def load_example_czi(**kwargs) -> 'CziScanner':
"""load example czi file"""
from neuralib.scan.czi import czi_file # pyright: ignore[reportMissingImports]
with google_drive_file('1gSPz_a7kCZ3UQABC-v-_sNnAxhfY_ly0', rename_file='test.czi', **kwargs) as file:
with czi_file(file) as czi:
return czi