diff --git a/README.md b/README.md index 21faa84..a515f13 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,18 @@ sudo cp cuda/lib64/libcudnn* /usr/local/cuda-10.0/lib64 sudo chmod a+r /usr/local/cuda-10.0/include/cudnn.h /usr/local/cuda-10.0/lib64/libcudnn* ``` +(optional): check CUDNN installation or for troubleshooting only) +Download and unzip https://ibl.flatironinstitute.org/resources/cudnn_samples_v7.zip +If necessary, setup your CUDA environment variables with the version you want to test + +``` +cd cudnn_samples_v7/mnistCUDNN/ +make clean && make +./mnistCUDNN +``` + +Should print a message that finishes with Test passed ! + ### Create a Python environment with TensorFlow and DLC Install a few things system wide and then python3.7 @@ -76,8 +88,10 @@ source ~/Documents/PYTHON/envs/dlcenv/bin/activate Install packages (please observe order of those commands!) ```bash -pip install "dask[complete]" -pip install ibllib +# pip install "dask[complete]" +pip install -U setuptools +# pip install git+https://github.com/int-brain-lab/ibllib.git@develop +pip install git+https://github.com/int-brain-lab/ibllib.git pip install https://extras.wxpython.org/wxPython4/extras/linux/gtk3/ubuntu-18.04/wxPython-4.0.7-cp37-cp37m-linux_x86_64.whl pip install tensorflow-gpu==1.13.1 pip install deeplabcut @@ -109,17 +123,27 @@ python -c 'import iblvideo' ``` ## Releasing a new version (for devs) -We use semantic versioning MAJOR.MINOR.PATCH. If you update the version, see below for what to adapt. Afterwards, tag the new version on Github. +We use semantic versioning MAJOR.MINOR.PATCH. If you update the version, see below for what to adapt. ### Any version update Update the version in ``` iblvideo/iblvideo/__init__.py ``` +Afterwards, tag the new version on Github. + ### Update MINOR or MAJOR -The version of weights and test data are synchronized with the MAJOR.MINOR version of this code. In addition to updating the version you have to upload two new zip files to FlatIron (note that the patch version is not included in the name): +The version of DLC weights and DLC test data are synchronized with the MAJOR.MINOR version of this code. (Note that the patch version is not included in the directory names) + +If you update any of the DLC weights, you also need to update the MINOR version of the code and the DLC test data, and vice versa. +1. For the weights, create a new directory called `weights_v{MAJOR}.{MINOR}` and copy the new weights, plus any unchanged weights into it. +2. Make a new `dlc_test_data_v{MAJOR}.{MINOR}` directory, with subdirectories `input` and `output`. +3. Copy the three videos from the `input` folder of the previous version dlc_test_data to the new one. +4. Create the three parquet files to go in `output` by running iblvideo.dlc() with the new weights folder as `path_dlc`, and each of the videos in the new `input` folder as `file_mp4`. +5. Zip and upload the new weights and test data folders to FlatIron : ``` /resources/dlc/weights_v{MAJOR}.{MINOR}.zip -/integration/dlc/test_data/test_data_v{MAJOR}.{MINOR}.zip +/integration/dlc/test_data/dlc_test_data_v{MAJOR}.{MINOR}.zip ``` +6. Delete your local weights and test data and run tests/test_choiceworld.py to make sure everything worked. diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..fd7db97 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,27 @@ +#docker build -t ibl/dlc:base . +FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 + +# link the cuda libraries +ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH + +# setup time zone for tz +ENV TZ=Europe/Paris +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone + +# Install python3.7 +RUN apt-get update +RUN apt-get install -y software-properties-common +RUN add-apt-repository ppa:deadsnakes/ppa +RUN apt-get update +RUN apt-get install -y python3.7 python3.7-dev python3.7-tk python3-pip python3.7-venv git ffmpeg libgtk-3-dev + +# Install Python dependencies +ARG PYTHON=python3.7 +ENV LANG C.UTF-8 +RUN ln -sf /usr/bin/${PYTHON} /usr/local/bin/python3 +RUN python3 -m pip install --upgrade pip +RUN python3 -m pip install "dask[complete]" +RUN python3 -m pip install ibllib +RUN python3 -m pip install https://extras.wxpython.org/wxPython4/extras/linux/gtk3/ubuntu-18.04/wxPython-4.0.7-cp37-cp37m-linux_x86_64.whl +RUN python3 -m pip install tensorflow-gpu==1.13.1 +RUN python3 -m pip install deeplabcut diff --git a/iblvideo/__init__.py b/iblvideo/__init__.py index 26e5dd5..58739e1 100644 --- a/iblvideo/__init__.py +++ b/iblvideo/__init__.py @@ -1,9 +1,9 @@ -__version__ = '0.2.0' # This is the only place where the version is hard coded, only adapt here +__version__ = '1.0.0' # This is the only place where the version is hard coded, only adapt here import deeplabcut from iblvideo.run import run_session, run_queue -from iblvideo.choiceworld import dlc, dlc_parallel +from iblvideo.choiceworld import dlc from iblvideo.weights import download_weights from iblvideo.cluster import create_cpu_gpu_cluster diff --git a/iblvideo/choiceworld.py b/iblvideo/choiceworld.py index f59b33b..7d7e088 100644 --- a/iblvideo/choiceworld.py +++ b/iblvideo/choiceworld.py @@ -1,15 +1,18 @@ """Functions to run DLC on IBL data with existing networks.""" -import deeplabcut +import deeplabcut # needs to be imported first import os -from glob import glob import shutil import logging import yaml +import time +from glob import glob from pathlib import Path +from collections import OrderedDict + import numpy as np import pandas as pd import cv2 -import time + from iblvideo.params import BODY_FEATURES, SIDE_FEATURES, LEFT_VIDEO, RIGHT_VIDEO, BODY_VIDEO from iblvideo.cluster import create_cpu_gpu_cluster from iblvideo.utils import _run_command @@ -61,16 +64,16 @@ def _dlc_init(file_mp4, path_dlc): return file_mp4, dlc_params, networks, tdir, tfile, file_label -def _get_crop_window(df_crop, network): +def _get_crop_window(file_df_crop, network): """ Get average position of a pivot point for autocropping. - - :param df_crop: data frame from hdf5 file from video data + :param file_df_crop: Path to data frame from hdf5 file from video data :param network: dictionary describing the networks. See constants SIDE and BODY :return: list of floats [width, height, x, y] defining window used for ffmpeg crop command """ + df_crop = pd.read_hdf(file_df_crop) XYs = [] for part in network['features']: x_values = df_crop[(df_crop.keys()[0][0], part, 'x')].values @@ -85,163 +88,198 @@ def _get_crop_window(df_crop, network): XYs.append([int(np.nanmean(x)), int(np.nanmean(y))]) xy = np.mean(XYs, axis=0) - return network['crop'](*xy) -def _s00_transform_rightCam(file_mp4, tdir): +def _s00_transform_rightCam(file_mp4, tdir, force=False): """ Flip and rotate the right cam and increase spatial resolution. - Such that the rightCamera video looks like the leftCamera video. """ - # TODO use the video parameters above not to have this hard-coded - # (sampling + original size) - if 'rightCamera' not in file_mp4.name: - return file_mp4 - + file_out1 = str(Path(tdir).joinpath(str(file_mp4).replace('.raw.', '.flipped.'))) + # If flipped right cam does not exist, compute + if os.path.exists(file_out1) and force is not True: + _logger.info('STEP 00a Flipped rightCamera video exists, not computing.') else: - _logger.info('STEP 00 Flipping and turning rightCamera video') - file_out1 = str(Path(tdir).joinpath(str(file_mp4).replace('.raw.', '.flipped.'))) + _logger.info('STEP 00a START Flipping and turning rightCamera video') command_flip = (f'ffmpeg -nostats -y -loglevel 0 -i {file_mp4} -vf ' f'"transpose=1,transpose=1" -vf hflip {file_out1}') pop = _run_command(command_flip) if pop['process'].returncode != 0: _logger.error(f' DLC 0a/5: Flipping ffmpeg failed: {file_mp4} ' + pop['stderr']) - _logger.info('Oversampling rightCamera video') - file_out2 = file_out1.replace('.flipped.', '.raw.transformed.') - + _logger.info('STEP 00a END Flipping and turning rightCamera video') + # Set force to true to recompute all subsequent steps + force = True + + # If oversampled cam does not exist, compute + file_out2 = file_out1.replace('.flipped.', '.raw.transformed.') + if os.path.exists(file_out2) and force is not True: + _logger.info('STEP 00b Oversampled rightCamera video exists, not computing.') + else: + _logger.info('STEP 00b START Oversampling rightCamera video') command_upsample = (f'ffmpeg -nostats -y -loglevel 0 -i {file_out1} ' f'-vf scale=1280:1024 {file_out2}') pop = _run_command(command_upsample) if pop['process'].returncode != 0: _logger.error(f' DLC 0b/5: Increase reso ffmpeg failed: {file_mp4}' + pop['stderr']) - Path(file_out1).unlink() - _logger.info('STEP 00 STOP Flipping and turning rightCamera video') + _logger.info('STEP 00b END Oversampling rightCamera video') + # Set force to true to recompute all subsequent steps + force = True - return file_out2 + return file_out2, force def _s01_subsample(file_in, file_out, force=False): """ Step 1 subsample video for detection. - Put 500 uniformly sampled frames into new video. """ - _logger.info(f"STEP 01 Generating sparse frame video {file_out} for posture detection") file_in = Path(file_in) file_out = Path(file_out) - if file_out.exists() and not force: - return file_out - - cap = cv2.VideoCapture(str(file_in)) - frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - - # get from 20 to 500 samples linearly spaced throughout the session - nsamples = min(max(20, frameCount / cap.get(cv2.CAP_PROP_FPS)), 500) - samples = np.int32(np.round(np.linspace(0, frameCount - 1, nsamples))) - - size = (int(cap.get(3)), int(cap.get(4))) - out = cv2.VideoWriter(str(file_out), cv2.VideoWriter_fourcc(*'mp4v'), 5, size) - for i in samples: - cap.set(1, i) - _, frame = cap.read() - out.write(frame) - - out.release() - _logger.info(f"STEP 01 STOP Generating sparse frame video {file_out} for posture detection") - return file_out - - -def _s02_detect_rois(tpath, sparse_video, dlc_params, create_labels=False): + if file_out.exists() and force is not True: + _logger.info(f"STEP 01 Sparse frame video {file_out.name} exists, not computing") + else: + _logger.info(f"STEP 01 START Generating sparse video {file_out.name} for posture" + f" detection") + cap = cv2.VideoCapture(str(file_in)) + frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # get from 20 to 500 samples linearly spaced throughout the session + nsamples = min(max(20, frameCount / cap.get(cv2.CAP_PROP_FPS)), 500) + samples = np.int32(np.round(np.linspace(0, frameCount - 1, nsamples))) + + size = (int(cap.get(3)), int(cap.get(4))) + out = cv2.VideoWriter(str(file_out), cv2.VideoWriter_fourcc(*'mp4v'), 5, size) + for i in samples: + cap.set(1, i) + _, frame = cap.read() + out.write(frame) + out.release() + _logger.info(f"STEP 01 END Generating sparse video {file_out.name} for posture detection") + # Set force to true to recompute all subsequent steps + force = True + + return file_out, force + + +def _s02_detect_rois(tdir, sparse_video, dlc_params, create_labels=False, force=False): """ Step 2 run DLC to detect ROIS. - - returns: df_crop, dataframe used to crop video + returns: Path to dataframe used to crop video """ - _logger.info(f"STEP 02 START Posture detection {sparse_video}") - out = deeplabcut.analyze_videos(dlc_params['roi_detect'], [str(sparse_video)]) - if create_labels: - deeplabcut.create_labeled_video(dlc_params['roi_detect'], [str(sparse_video)]) - h5_sub = next(tpath.glob(f'*{out}*.h5'), None) - _logger.info(f"STEP 02 END Posture detection {sparse_video}") - file_out = pd.read_hdf(h5_sub) - return file_out + file_out = next(tdir.glob('*subsampled*.h5'), None) + if file_out is None or force is True: + _logger.info(f"STEP 02 START Posture detection for {sparse_video.name}") + out = deeplabcut.analyze_videos(dlc_params['roi_detect'], [str(sparse_video)]) + if create_labels: + deeplabcut.create_labeled_video(dlc_params['roi_detect'], [str(sparse_video)]) + file_out = next(tdir.glob(f'*{out}*.h5'), None) + _logger.info(f"STEP 02 END Posture detection for {sparse_video.name}") + # Set force to true to recompute all subsequent steps + force = True + else: + _logger.info(f"STEP 02 Posture detection for {sparse_video.name} exists, not computing.") + return file_out, force -def _s03_crop_videos(df_crop, file_in, file_out, network): +def _s03_crop_videos(file_df_crop, file_in, file_out, network, force=False): """ Step 3 crop videos using ffmpeg. - returns: dictionary of cropping coordinates relative to upper left corner """ - _logger.info(f'STEP 03 START cropping {network["label"]} video') - crop_command = ('ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf "crop={w[0]}:{w[1]}:' - '{w[2]}:{w[3]}" -c:v libx264 -crf 11 -c:a copy {file_out}') - whxy = _get_crop_window(df_crop, network) - pop = _run_command(crop_command.format(file_in=file_in, file_out=file_out, w=whxy)) - if pop['process'].returncode != 0: - _logger.error(f'DLC 3/6: Cropping ffmpeg failed for ROI \ - {network["label"]}, file: {file_in}') - np.save(file_out.parent.joinpath(file_out.stem + '.whxy.npy'), whxy) - _logger.info(f'STEP 03 STOP cropping {network["label"]} video') - return file_out - - -def _s04_brightness_eye(file_mp4, force=False): + # Don't run if outputs exist and force is False + file_out = Path(file_out) + whxy_file = file_out.parent.joinpath(file_out.stem + '.whxy.npy') + if file_out.exists() and whxy_file.exists() and force is not True: + _logger.info(f'STEP 03 Cropped video {file_out.name} exists, not computing.') + else: + _logger.info(f'STEP 03 START generating cropped video {file_out.name}') + crop_command = ('ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf "crop={w[0]}:{w[1]}:' + '{w[2]}:{w[3]}" -c:v libx264 -crf 11 -c:a copy {file_out}') + whxy = _get_crop_window(file_df_crop, network) + pop = _run_command(crop_command.format(file_in=file_in, file_out=file_out, w=whxy)) + if pop['process'].returncode != 0: + _logger.error(f'DLC 3/6: Cropping ffmpeg failed for ROI \ + {network["label"]}, file: {file_in}') + np.save(str(whxy_file), whxy) + _logger.info(f'STEP 03 END generating cropped video {file_out.name}') + # Set force to true to recompute all subsequent steps + force = True + return file_out, force + + +def _s04_brightness_eye(file_in, force=False): """ Adjust brightness for eye for better network performance. - wget -O- http://ffmpeg.org/releases/ffmpeg-snapshot.tar.bz2 | tar xj """ - file_out = file_mp4 - file_in = file_mp4.parent.joinpath(file_mp4.name.replace('eye', 'eye.nobright')) - if file_in.exists() and not force: - return file_out - _logger.info('STEP 04 START Adjusting eye brightness') - file_out.rename(file_in) - cmd = (f'ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf ' - f'colorlevels=rimax=0.25:gimax=0.25:bimax=0.25 -c:a copy {file_out}') - pop = _run_command(cmd) - if pop['process'].returncode != 0: - _logger.error(f"DLC 4/6: (str(dlc_params), [str(tfile)]) failed: {file_in}") - _logger.info('STEP 04 STOP Adjusting eye brightness') - return file_out - - -def _s04_resample_paws(file_mp4, tdir, force=False): + # This function renames the input to 'eye.nobright' and then saves the adjusted + # output under the same name as the original input. Therefore: + file_in = Path(file_in) + file_out = file_in.parent.joinpath(file_in.name.replace('eye', 'eye_adjusted')) + + if file_out.exists() and force is not True: + _logger.info(f'STEP 04 Adjusted eye brightness {file_out.name} exists, not computing') + else: + # Else run command + _logger.info(f'STEP 04 START Generating adjusting eye brightness video {file_out.name}') + cmd = (f'ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf ' + f'colorlevels=rimax=0.25:gimax=0.25:bimax=0.25 -c:a copy {file_out}') + pop = _run_command(cmd) + if pop['process'].returncode != 0: + _logger.error(f"DLC 4/6: Adjust eye brightness failed: {file_in}") + _logger.info(f'STEP 04 END Generating adjusting eye brightness video {file_out.name}') + # Set force to true to recompute all subsequent steps + force = True + return file_out, force + + +def _s04_resample_paws(file_in, tdir, force=False): """For paws, spatially downsample to speed up processing x100.""" - file_mp4 = Path(file_mp4) - file_in = file_mp4 - file_out = Path(tdir) / file_mp4.name.replace('raw', 'paws_downsampled') + file_in = Path(file_in) + file_out = Path(tdir) / file_in.name.replace('raw', 'paws_downsampled') - cmd = (f'ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf scale=128:102 -c:v libx264 -crf 23' - f' -c:a copy {file_out}') # was 112:100 - pop = _run_command(cmd) - if pop['process'].returncode != 0: - _logger.error(f"DLC 4/6: Subsampling paws failed: {file_in}") - _logger.info('STEP 04 STOP resample paws') - return file_out + if file_out.exists() and force is not True: + _logger.info(f'STEP 04 resampled paws {file_out.name} exists, not computing') + else: + _logger.info(f'STEP 04 START generating resampled paws video {file_out.name}') + cmd = (f'ffmpeg -nostats -y -loglevel 0 -i {file_in} -vf scale=128:102 -c:v libx264 -crf 23' + f' -c:a copy {file_out}') # was 112:100 + pop = _run_command(cmd) + if pop['process'].returncode != 0: + _logger.error(f"DLC 4/6: Subsampling paws failed: {file_in}") + _logger.info(f'STEP 04 END generating resampled paws video {file_out.name}') + # Set force to true to recompute all subsequent steps + force = True + return file_out, force -def _s05_run_dlc_specialized_networks(dlc_params, tfile, network, create_labels=False): - _logger.info(f'STEP 05 START extract dlc feature {tfile}') - deeplabcut.analyze_videos(str(dlc_params), [str(tfile)]) - if create_labels: - deeplabcut.create_labeled_video(str(dlc_params), [str(tfile)]) - deeplabcut.filterpredictions(str(dlc_params), [str(tfile)]) - _logger.info(f'STEP 05 STOP extract dlc feature {tfile}') - return network +def _s05_run_dlc_specialized_networks(dlc_params, tfile, network, create_labels=False, + force=False): + + # Check if final result exists + result = next(tfile.parent.glob(f'*{network}*filtered.h5'), None) + if result and force is not True: + _logger.info(f'STEP 05 dlc feature for {tfile.name} already extracted, not computing.') + else: + _logger.info(f'STEP 05 START extract dlc feature for {tfile.name}') + deeplabcut.analyze_videos(str(dlc_params), [str(tfile)]) + if create_labels: + deeplabcut.create_labeled_video(str(dlc_params), [str(tfile)]) + deeplabcut.filterpredictions(str(dlc_params), [str(tfile)]) + _logger.info(f'STEP 05 END extract dlc feature for {tfile.name}') + # Set force to true to recompute all subsequent steps + force = True + return def _s06_extract_dlc_alf(tdir, file_label, networks, file_mp4, *args): """ Output an ALF matrix. - Column names contain the full DLC results [nframes, nfeatures] """ - _logger.info('STEP 06 START wrap-up and extract ALF files') + _logger.info(f'STEP 06 START wrap-up and extract ALF files {file_label}') if 'bodyCamera' in file_label: video_params = BODY_VIDEO elif 'leftCamera' in file_label: @@ -289,11 +327,11 @@ def _s06_extract_dlc_alf(tdir, file_label, networks, file_mp4, *args): file_alf = alf_path.joinpath(f'_ibl_{file_label}.dlc.pqt') df_full.to_parquet(file_alf) - _logger.info('STEP 06 STOP wrap-up and extract ALF files') + _logger.info(f'STEP 06 END wrap-up and extract ALF files {file_label}') return file_alf -def dlc(file_mp4, path_dlc=None, force=False): +def dlc(file_mp4, path_dlc=None, force=False, dlc_timer=None): """ Analyse a leftCamera, rightCamera or bodyCamera video with DeepLabCut. @@ -308,110 +346,140 @@ def dlc(file_mp4, path_dlc=None, force=False): :param file_mp4: Video file to run :param path_dlc: Path to folder with DLC weights + :param force: bool, whether to overwrite existing intermediate files :return out_file: Path to DLC table in parquet file format """ - start_T = time.time() + dlc_timer = dlc_timer or OrderedDict() + time_total_on = time.time() # Initiate file_mp4, dlc_params, networks, tdir, tfile, file_label = _dlc_init(file_mp4, path_dlc) # Run the processing steps in order - file2segment = _s00_transform_rightCam(file_mp4, tdir) # CPU pure Python - file_sparse = _s01_subsample(file2segment, tfile['mp4_sub']) # CPU ffmpeg - df_crop = _s02_detect_rois(tdir, file_sparse, dlc_params) # GPU dlc - - networks_run = {} + if 'rightCamera' not in file_mp4.name: + file2segment = file_mp4 + else: + time_on = time.time() + file2segment, force = _s00_transform_rightCam(file_mp4, tdir, force=force) # CPU Python + time_off = time.time() + dlc_timer['Transform right camera'] = time_off - time_on + + time_on = time.time() + file_sparse, force = _s01_subsample(file2segment, tfile['mp4_sub'], force=force) # CPU ffmpeg + time_off = time.time() + dlc_timer['Subsample video'] = time_off - time_on + + time_on = time.time() + file_df_crop, force = _s02_detect_rois(tdir, file_sparse, dlc_params, force=force) # GPU dlc + time_off = time.time() + dlc_timer['Detect ROIs'] = time_off - time_on + + input_force = force for k in networks: + time_on = time.time() if networks[k]['features'] is None: continue + # Run preprocessing depdening on the feature if k == 'paws': - cropped_vid = _s04_resample_paws(file2segment, tdir) + preproc_vid, force = _s04_resample_paws(file2segment, tdir, force=force) elif k == 'eye': - cropped_vid_a = _s03_crop_videos(df_crop, file2segment, tfile[k], networks[k]) - cropped_vid = _s04_brightness_eye(cropped_vid_a) + cropped_vid, force = _s03_crop_videos(file_df_crop, file2segment, tfile[k], + networks[k], force=force) + preproc_vid, force = _s04_brightness_eye(cropped_vid, force=force) else: - cropped_vid = _s03_crop_videos(df_crop, file2segment, tfile[k], networks[k]) - network_run = _s05_run_dlc_specialized_networks(dlc_params[k], cropped_vid, networks[k]) - networks_run[k] = network_run + preproc_vid, force = _s03_crop_videos(file_df_crop, file2segment, tfile[k], + networks[k], force=force) + time_off = time.time() + dlc_timer[f'Prepare video for {k} network'] = time_off - time_on - out_file = _s06_extract_dlc_alf(tdir, file_label, networks_run, file_mp4) + time_on = time.time() + _s05_run_dlc_specialized_networks(dlc_params[k], preproc_vid, k, force=force) + time_off = time.time() + dlc_timer[f'Run {k} network'] = time_off - time_on + # Reset force to the original input value as the reset is network-specific + force = input_force - file2segment = Path(file2segment) - # at the end mop up the mess - shutil.rmtree(tdir) + out_file = _s06_extract_dlc_alf(tdir, file_label, networks, file_mp4) + # at the end mop up the mess + # For right camera video only + file2segment = Path(file2segment) if '.raw.transformed' in file2segment.name: file2segment.unlink() - - # Back to home folder else there are conflicts in a loop - os.chdir(Path.home()) - end_T = time.time() - print(file_label) - print('In total this took: ', end_T - start_T) - - return out_file - - -def dlc_parallel(file_mp4, path_dlc=None, force=False): - """ - Run dlc in parallel. - - :param file_mp4: Video file to run - :param path_dlc: Path to folder with DLC weights - :return out_file: Path to DLC table in parquet file format - """ - - import dask - cluster, client = create_cpu_gpu_cluster() - - start_T = time.time() - file_mp4, dlc_params, networks, tdir, tfile, file_label = _dlc_init(file_mp4, path_dlc) - - # Run the processing steps in order - future_s00 = client.submit(_s00_transform_rightCam, file_mp4, tdir, workers='CPU') - file2segment = future_s00.result() - - future_s01 = client.submit(_s01_subsample, file2segment, tfile['mp4_sub'], workers='CPU') - file_sparse = future_s01.result() - - future_s02 = client.submit(_s02_detect_rois, tdir, file_sparse, dlc_params, workers='GPU') - df_crop = future_s02.result() - - networks_run = {} - for k in networks: - if networks[k]['features'] is None: - continue - if k == 'paws': - future_s04 = client.submit(_s04_resample_paws, file2segment, tdir, workers='CPU') - cropped_vid = future_s04.result() - elif k == 'eye': - future_s03 = client.submit(_s03_crop_videos, df_crop, file2segment, tfile[k], - networks[k], workers='CPU') - cropped_vid_a = future_s03.result() - future_s04 = client.submit(_s04_brightness_eye, cropped_vid_a, workers='CPU') - cropped_vid = future_s04.result() - else: - future_s03 = client.submit(_s03_crop_videos, df_crop, file2segment, tfile[k], - networks[k], workers='CPU') - cropped_vid = future_s03.result() - - future_s05 = client.submit(_s05_run_dlc_specialized_networks, dlc_params[k], cropped_vid, - networks[k], workers='GPU') - network_run = future_s05.result() - networks_run[k] = network_run - - pipeline = dask.delayed(_s06_extract_dlc_alf)(tdir, file_label, networks_run, file_mp4) - future = client.compute(pipeline) - out_file = future.result() - - cluster.close() - client.close() + flipped = Path(str(file2segment).replace('raw.transformed', 'flipped')) + flipped.unlink() shutil.rmtree(tdir) - # Back to home folder else there are conflicts in a loop os.chdir(Path.home()) end_T = time.time() print(file_label) - print('In total this took: ', end_T - start_T) - - return out_file + time_total_off = time.time() + dlc_timer['DLC total'] = time_total_off - time_total_on + + return out_file, dlc_timer + + +# def dlc_parallel(file_mp4, path_dlc=None, force=False): +# """ +# Run dlc in parallel. +# +# :param file_mp4: Video file to run +# :param path_dlc: Path to folder with DLC weights +# :return out_file: Path to DLC table in parquet file format +# """ +# +# import dask +# cluster, client = create_cpu_gpu_cluster() +# +# start_T = time.time() +# file_mp4, dlc_params, networks, tdir, tfile, file_label = _dlc_init(file_mp4, path_dlc) +# +# # Run the processing steps in order +# future_s00 = client.submit(_s00_transform_rightCam, file_mp4, tdir, workers='CPU') +# file2segment = future_s00.result() +# +# future_s01 = client.submit(_s01_subsample, file2segment, tfile['mp4_sub'], workers='CPU') +# file_sparse = future_s01.result() +# +# future_s02 = client.submit(_s02_detect_rois, tdir, file_sparse, dlc_params, workers='GPU') +# df_crop = future_s02.result() +# +# networks_run = {} +# for k in networks: +# if networks[k]['features'] is None: +# continue +# if k == 'paws': +# future_s04 = client.submit(_s04_resample_paws, file2segment, tdir, workers='CPU') +# cropped_vid = future_s04.result() +# elif k == 'eye': +# future_s03 = client.submit(_s03_crop_videos, df_crop, file2segment, tfile[k], +# networks[k], workers='CPU') +# cropped_vid_a = future_s03.result() +# future_s04 = client.submit(_s04_brightness_eye, cropped_vid_a, workers='CPU') +# cropped_vid = future_s04.result() +# else: +# future_s03 = client.submit(_s03_crop_videos, df_crop, file2segment, tfile[k], +# networks[k], workers='CPU') +# cropped_vid = future_s03.result() +# +# future_s05 = client.submit(_s05_run_dlc_specialized_networks, dlc_params[k], cropped_vid, +# networks[k], workers='GPU') +# network_run = future_s05.result() +# networks_run[k] = network_run +# +# pipeline = dask.delayed(_s06_extract_dlc_alf)(tdir, file_label, networks_run, file_mp4) +# future = client.compute(pipeline) +# out_file = future.result() +# +# cluster.close() +# client.close() +# +# shutil.rmtree(tdir) +# +# # Back to home folder else there are conflicts in a loop +# os.chdir(Path.home()) +# end_T = time.time() +# print(file_label) +# print('In total this took: ', end_T - start_T) +# +# return out_file diff --git a/iblvideo/motion_energy.py b/iblvideo/motion_energy.py index e430596..2a60483 100644 --- a/iblvideo/motion_energy.py +++ b/iblvideo/motion_energy.py @@ -1,4 +1,4 @@ -''' +""" For a session where there is DLC already computed, load DLC traces to cut video ROIs and then compute motion energy for these ROIS. @@ -6,66 +6,70 @@ bodyCamera: cut ROI such that mouse body but not wheel motion is in ROI left(right)Camera: cut whisker pad region -''' +""" import os import time -import logging -from pathlib import Path import numpy as np +import pandas as pd import cv2 from oneibl.one import ONE -import alf.io -from ibllib.io.video import get_video_frames_preload, url_from_eid -_logger = logging.getLogger('ibllib') +from ibllib.io.video import get_video_frames_preload, url_from_eid, label_from_path +from ibllib.io.extractors.camera import get_video_length +from oneibl.stream import VideoStreamer -def get_dlc_midpoints(eid, side, one=None): +def grayscale(x): + return cv2.cvtColor(x, cv2.COLOR_BGR2GRAY) - if one is None: - one = ONE() - - # Download dlc data if not available locally - _ = one.load(eid, dataset_types='camera.dlc', download_only=True) - local_path = one.path_from_eid(eid) - alf_path = Path(local_path).joinpath('alf') - - dlc = alf.io.load_object(alf_path, f'{side}Camera', namespace='ibl')['dlc'] +def get_dlc_midpoints(dlc_pqt): + # Load dataframe + dlc_df = pd.read_parquet(dlc_pqt) # Set values to nan if likelihood is too low and calcualte midpoints - targets = np.unique(['_'.join(col.split('_')[:-1]) for col in dlc.columns]) + targets = np.unique(['_'.join(col.split('_')[:-1]) for col in dlc_df.columns]) mloc = {} for t in targets: - idx = dlc.loc[dlc[f'{t}_likelihood'] < 0.9].index - dlc.loc[idx, [f'{t}_x', f'{t}_y']] = np.nan - mloc[t] = [int(np.nanmean(dlc[f'{t}_x'])), int(np.nanmean(dlc[f'{t}_y']))] - + idx = dlc_df.loc[dlc_df[f'{t}_likelihood'] < 0.9].index + dlc_df.loc[idx, [f'{t}_x', f'{t}_y']] = np.nan + mloc[t] = [int(np.nanmean(dlc_df[f'{t}_x'])), int(np.nanmean(dlc_df[f'{t}_y']))] return mloc -def compute_ROI_ME(eid, side, frame_numbers=None, one=None): - ''' +def motion_energy(session_path, dlc_pqt, frames=None, one=None): + """ Compute motion energy on cropped frames of a single video - :param eid: Session ID - :param side: 'body', 'left' or 'right' - ''' - - if one is None: - one = ONE() - + :param session_path: Path to session. + :param dlc_pqt: Path to dlc result in pqt file format. If None all frames are loaded at once. + :param frames: Number of frames to load into memory at once. + :param one: ONE instance + :return me_file: Path to numpy file contaiing motion energy. + :return me_roi: Path to numpy file containing ROI coordinates. + + The frames parameter determines how many cropped frames per camera are loaded into memory at + once and should be set depending on availble RAM. Some approximate numbers for orientation, + assuming 90 min video and frames set to: + 1 : 152 KB (body), 54 KB (left), 15 KB (right) + 50000 : 7.6 GB (body), 2.7 GB (left), 0.75 GB (right) + None : 25 GB (body), 17.5 GB (left), 12.5 GB (right) + """ + + one = one or ONE() start_T = time.time() - video_path = one.path_from_eid(eid).joinpath('raw_video_data', - f'_iblrig_{side}Camera.raw.mp4') + + # Get label from dlc_df + label = label_from_path(dlc_pqt) + video_path = session_path.joinpath('raw_video_data', f'_iblrig_{label}Camera.raw.mp4') # Check if video available locally, else create url if not os.path.isfile(video_path): - video_path = url_from_eid(eid, label=side, one=one) + eid = one.eid_from_path(session_path) + video_path = url_from_eid(eid, label=label, one=one) # Crop ROI - _logger.info('{side}Camera Cropping ROI') - mloc = get_dlc_midpoints(eid, side, one=one) - if side == 'body': + mloc = get_dlc_midpoints(dlc_pqt) + if label == 'body': anchor = np.array(mloc['tail_start']) w, h = int(anchor[0] * 3 / 5), 210 x, y = int(anchor[0] - anchor[0] * 3 / 5), int(anchor[1] - 120) @@ -78,26 +82,48 @@ def compute_ROI_ME(eid, side, frame_numbers=None, one=None): # Note that x and y are flipped when loading with cv2, therefore: mask = np.s_[y:y + h, x:x + w] + # save ROI coordinates + roi = np.asarray([w, h, x, y]) + alf_path = session_path.joinpath('alf') + roi_file = alf_path.joinpath(f'{label}ROIMotionEnergy.position.npy') + np.save(roi_file, roi) + + frame_count = get_video_length(video_path) + me = np.zeros(frame_count,) + + is_url = isinstance(video_path, str) and video_path.startswith('http') + cap = VideoStreamer(video_path).cap if is_url else cv2.VideoCapture(str(video_path)) + if frames: + n, keep_reading = 0, True + while keep_reading: + # Set the frame numbers to the next #frames, with 1 frame overlap + frame_numbers = range(n * (frames - 1), n * (frames - 1) + frames) + # Make sure not to load empty frames + if np.max(frame_numbers) >= frame_count: + frame_numbers = range(frame_numbers.start, frame_count) + keep_reading = False + # Load, crop and grayscale frames. + cropped_frames = get_video_frames_preload(cap, frame_numbers=frame_numbers, + mask=mask, func=grayscale, + quiet=True).astype(np.float32) + # Calculate motion energy for those frames and append to big array + me[frame_numbers[:-1]] = np.mean(np.abs(np.diff(cropped_frames, axis=0)), axis=(1, 2)) + # Next set of frames + n += 1 + else: + # Compute on entire video at once + cropped_frames = get_video_frames_preload(cap, frame_numbers=None, mask=mask, + func=grayscale, quiet=True).astype(np.float32) + me[:-1] = np.mean(np.abs(np.diff(cropped_frames, axis=0)), axis=(1, 2)) - # Crop and grayscale frames. - cropped_frames = get_video_frames_preload(video_path, frame_numbers=frame_numbers, mask=mask, - func=cv2.cvtColor, code=cv2.COLOR_BGR2GRAY) - - # save ROI - alf_path = one.path_from_eid(eid).joinpath('alf') - np.save(alf_path.joinpath(f'{side}ROIMotionEnergy.position.npy'), np.asarray([w, h, x, y])) - - # Compute and save motion energy - _logger.info(f'{side}Camera computing motion energy') - # Cast to float - cropped_frames = np.asarray(cropped_frames, dtype=np.float32) - me = np.mean(np.abs(cropped_frames[1:] - cropped_frames[:-1]), axis=(1, 2)) # copy last value to make motion energy fit frame length - me = np.append(me, me[-1]) + cap.release() + me[-1] = me[-2] # save ME - np.save(alf_path.joinpath(f'{side}.ROIMotionEnergy.npy'), me) - - _logger.info(f'Motion energy and ROI for {side}Camera computed and saved') + me_file = alf_path.joinpath(f'{label}Camera.ROIMotionEnergy.npy') + np.save(me_file, me) end_T = time.time() - print(f'{side}Camera computed in', np.round((end_T - start_T), 2)) + print(f'{label}Camera computed in', np.round((end_T - start_T), 2)) + + return me_file, roi_file diff --git a/iblvideo/run.py b/iblvideo/run.py index 32bb2bc..8aed100 100644 --- a/iblvideo/run.py +++ b/iblvideo/run.py @@ -1,86 +1,176 @@ import logging import shutil +import os import traceback +import time from datetime import datetime +from collections import OrderedDict + import numpy as np from oneibl.one import ONE from oneibl.patcher import FTPPatcher from iblvideo.choiceworld import dlc +from iblvideo.motion_energy import motion_energy from iblvideo.weights import download_weights from iblvideo import __version__ from ibllib.pipes import tasks +from ibllib.qc.dlc import DlcQC +from ibllib.io.video import assert_valid_label +from ibllib.exceptions import ALFObjectNotFound _logger = logging.getLogger('ibllib') # re-using the Task class allows to not re-write all the logging, error management # and automatic settings of task statuses in the database +def _format_timer(timer): + logstr = '' + for item in timer.items(): + logstr += f'\nTiming {item[0]}Camera [sec]\n' + for subitem in item[1].items(): + logstr += f'{subitem[0]}: {int(np.round(subitem[1]))}\n' + return logstr + + class TaskDLC(tasks.Task): gpu = 1 cpu = 4 io_charge = 90 level = 0 - def _run(self, machine=None, version=__version__): - # Download weights into ONE Cache directory under 'resources/DLC' if not exist - if machine is not None: - _logger.info(f'Running on {machine}') - path_dlc = download_weights(version=version) - files_mp4 = list(self.session_path.joinpath('raw_video_data').glob('*.mp4')) - _logger.info(f'Running DLC on {len(files_mp4)} video files.') - # Run dlc on all videos - out_files = [] - for cam in range(len(files_mp4)): - dlc_result = dlc(files_mp4[cam], path_dlc) - out_files.append(dlc_result) - _logger.info(dlc_result) - pqts = list(self.session_path.joinpath('alf').glob('*.pqt')) - return pqts - - -def run_session(session_id, machine=None, n_cams=3, one=None, version=__version__, - remove_data=True): + def _result_exists(self, session_id, fname): + """ Checks if dlc result is available locally or in database. """ + result = None + if os.path.exists(self.session_path.joinpath('alf', fname)): + result = self.session_path.joinpath('alf', fname) + _logger.info(f'Using local version of {fname}') + else: + try: + result = self.one.load_dataset(session_id, dataset=fname, download_only=True) + _logger.info(f'Downloaded {fname} from database') + except ALFObjectNotFound: + pass + return result + + def _run(self, cams=('left', 'body', 'right'), version=__version__, frames=None, **kwargs): + session_id = self.one.eid_from_path(self.session_path) + overwrite = kwargs.pop('overwrite', None) + # Create dictionary for logging time spent on each task + timer = OrderedDict() + # Loop through cams + dlc_results, me_results, me_rois = [], [], [] + for cam in cams: + timer[f'{cam}'] = OrderedDict() + # Check if dlc and me results are available locally or in database, if latter download + if overwrite: + # If it's a rerun, pretend the data doesn't exist yet + dlc_result, me_result, me_roi = None, None, None + else: + dlc_result = self._result_exists(session_id, f'_ibl_{cam}Camera.dlc.pqt') + me_result = self._result_exists(session_id, f'{cam}Camera.ROIMotionEnergy.npy') + me_roi = self._result_exists(session_id, f'{cam}ROIMotionEnergy.position.npy') + + # If dlc_result doesn't exist or should be overwritten, run DLC + if dlc_result is None: + # Download the camera data if not available locally + time_on = time.time() + _logger.info(f'Downloading {cam}Camera.') + file_mp4 = self.one.load_dataset(session_id, dataset=f'_iblrig_{cam}Camera.raw', + download_only=True) + time_off = time.time() + timer[f'{cam}'][f'Download video'] = time_off - time_on + # Download weights if not exist locally + time_on = time.time() + path_dlc = download_weights(version=version) + time_off = time.time() + timer[f'{cam}']['Download DLC weights'] = time_off - time_on + _logger.info(f'Running DLC on {cam}Camera.') + try: + dlc_result, timer[f'{cam}'] = dlc(file_mp4, path_dlc=path_dlc, force=overwrite, + dlc_timer=timer[f'{cam}']) + _logger.info(dlc_result) + except BaseException: + _logger.error(f'DLC {cam}Camera failed.\n' + traceback.format_exc()) + continue + dlc_results.append(dlc_result) + + # If me_results don't exist or should be overwritten, run me + if me_result is None or me_roi is None: + _logger.info(f'Computing motion energy for {cam}Camera') + try: + time_on = time.time() + me_result, me_roi = motion_energy(self.session_path, dlc_result, frames=frames, + one=self.one) + time_off = time.time() + timer[f'{cam}']['Compute motion energy'] = time_off - time_on + _logger.info(me_result) + _logger.info(me_roi) + except BaseException: + _logger.error(f'Motion energy {cam}Camera failed.\n' + traceback.format_exc()) + continue + me_results.append(me_result) + me_rois.append(me_roi) + _logger.info(_format_timer(timer)) + return dlc_results, me_results, me_rois + + +def run_session(session_id, machine=None, cams=('left', 'body', 'right'), one=None, + version=__version__, remove_videos=True, frames=50000, **kwargs): """ Run DLC on a single session in the database. :param session_id: Alyx eID of session to run :param machine: Tag for the machine this job ran on (string) - :param n_cams: Minimum number of camera datasets required + :param cams: Tuple of labels of videos to run dlc and motion energy on. + Valid labels are 'left', 'body' and 'right'. :param one: ONE instance to use for query (optional) :param version: Version of iblvideo / DLC weights to use (default is current version) - :param remove_data: Whether to remove the local raw_video_data after DLC (default is True) + :param remove_videos: Whether to remove the local raw_video_data after DLC (default is True) + :param frames: Number of video frames loaded into memory at once while computing ME. If None, + all frames of a video are loaded at once. (default is 50000, see below) + :param kwargs: Additional keyword arguments to be passed to TaskDLC :return status: final status of the task + + The frames parameter determines how many cropped frames per camera are loaded into memory at + once and should be set depending on availble RAM. Some approximate numbers for orientation, + assuming 90 min video and frames set to: + 1 : 152 KB (body), 54 KB (left), 15 KB (right) + 50000 : 7.6 GB (body), 2.7 GB (left), 0.75 GB (right) + None : 25 GB (body), 17.5 GB (left), 12.5 GB (right) """ # Catch all errors that are not caught inside run function and put them in the log try: - # Create ONE instance if none are given - if one is None: - one = ONE() - - # Find task for session and set to running + # Create ONE instance if none is given + one = one or ONE() session_path = one.path_from_eid(session_id) tdict = one.alyx.rest('tasks', 'list', django=f"name__icontains,DLC,session__id,{session_id}")[0] - one.alyx.rest('tasks', 'partial_update', id=tdict['id'], data={'status': 'Started'}) - - # Before starting to download, check if required number of cameras are available - dsets = one.alyx.rest('datasets', 'list', django=(f'session__id,{session_id},' - 'data_format__name,mp4,' - 'name__icontains,camera')) + except IndexError: + print(f"No DLC task found for session {session_id}") + return -1 - if len(dsets) < n_cams: + try: + # Check if labels are valid + cams = tuple(assert_valid_label(cam) for cam in cams) # raises ValueError if label invalid + # Check if all requested videos exist + vids = [dset['name'] for dset in one.alyx.rest('datasets', 'list', + django=(f'session__id,{session_id},' + 'data_format__name,mp4,' + f'name__icontains,camera') + )] + no_vid = [cam for cam in cams if f'_iblrig_{cam}Camera.raw.mp4' not in vids] + if len(no_vid) > 0: # If less datasets, update task and raise error - patch_data = {'log': f"Found only {len(dsets)} video files, user required {n_cams}.", - 'version': version, 'status': 'Errored'} + log_str = '\n'.join([f"No raw video file found for {no_cam}Camera." for + no_cam in no_vid]) + patch_data = {'log': log_str, 'status': 'Errored'} one.alyx.rest('tasks', 'partial_update', id=tdict['id'], data=patch_data) - status = -1 + return -1 else: - # Download camera files - one.load(session_id, dataset_types=['_iblrig_Camera.raw'], download_only=True) # create the task instance and run it, update task - task = TaskDLC(session_path, one=one, taskid=tdict['id']) - status = task.run(machine=machine, version=version) + task = TaskDLC(session_path, one=one, taskid=tdict['id'], machine=machine, **kwargs) + status = task.run(cams=cams, version=version, frames=frames) patch_data = {'time_elapsed_secs': task.time_elapsed_secs, 'log': task.log, 'version': version, 'status': 'Errored' if status == -1 else 'Complete'} one.alyx.rest('tasks', 'partial_update', id=tdict['id'], data=patch_data) @@ -88,36 +178,41 @@ def run_session(session_id, machine=None, n_cams=3, one=None, version=__version_ if task.outputs: # it is safer to instantiate the FTP right before transfer to prevent time-out ftp_patcher = FTPPatcher(one=one) - ftp_patcher.create_dataset(path=task.outputs) - - if remove_data is True: + ftp_patcher.create_dataset(path=task.outputs[0]) + ftp_patcher.create_dataset(path=task.outputs[1]) + ftp_patcher.create_dataset(path=task.outputs[2]) + if status == 0 and remove_videos is True: shutil.rmtree(session_path.joinpath('raw_video_data'), ignore_errors=True) + # Run DLC QC + # Download camera times and then force qc to use local data as dlc might not have + # been updated on FlatIron at this stage + one.load(session_id, dataset_types=['camera.times'], download_only=True) + for cam in cams: + qc = DlcQC(session_id, cam, one=one, download_data=False) + qc.run(update=True) + except BaseException: - patch_data = {'log': traceback.format_exc(), - 'version': version, 'status': 'Errored'} + patch_data = {'log': tdict['log'] + '\n\n' + traceback.format_exc(), 'status': 'Errored'} one.alyx.rest('tasks', 'partial_update', id=tdict['id'], data=patch_data) status = -1 return status -def run_queue(machine=None, n_sessions=np.inf, version=__version__, delta_query=600): +def run_queue(machine=None, n_sessions=np.inf, delta_query=600, **kwargs): """ Run the entire queue, or n_sessions, of DLC tasks on Alyx. :param machine: Tag for the machine this job ran on (string) :param n_sessions: Number of sessions to run from queue (default is run whole queue) - :param version: Version of iblvideo / DLC weights to use (default is current version) :param delta_query: Time between querying the database for Empty tasks, in sec + :param kwargs: Keyword arguments to be passed to run_session. """ - # Create ONE instance one = ONE() - - status_dict = {} - # Loop until n_sessions is reached or something breaks + status_dict = {} count = 0 last_query = datetime.now() while count < n_sessions: @@ -127,15 +222,13 @@ def run_queue(machine=None, n_sessions=np.inf, version=__version__, delta_query= last_query = datetime.now() tasks = one.alyx.rest('tasks', 'list', status='Empty', name='EphysDLC') sessions = [t['session'] for t in tasks] - # Return if no more sessions to run if len(sessions) == 0: print("No sessions to run") return - # Run next session in the list, capture status in dict and move on to next session eid = sessions[0] - status_dict[eid] = run_session(sessions.pop(0), machine=machine, one=one, version=version) + status_dict[eid] = run_session(sessions.pop(0), machine=machine, one=one, **kwargs) count += 1 return status_dict diff --git a/iblvideo/tests/__init__.py b/iblvideo/tests/__init__.py index 95a59c1..5a9e92a 100644 --- a/iblvideo/tests/__init__.py +++ b/iblvideo/tests/__init__.py @@ -1 +1 @@ -from iblvideo.tests.download_test_data import _download_test_data +from iblvideo.tests.download_test_data import _download_dlc_test_data, _download_me_test_data diff --git a/iblvideo/tests/download_test_data.py b/iblvideo/tests/download_test_data.py index efd4ba9..e9eddb3 100644 --- a/iblvideo/tests/download_test_data.py +++ b/iblvideo/tests/download_test_data.py @@ -2,10 +2,10 @@ from pathlib import Path from ibllib.io import params from oneibl.webclient import http_download_file -from .. import __version__ +from iblvideo import __version__ -def _download_test_data(version=__version__): +def _download_dlc_test_data(version=__version__,): """Download test data from FlatIron.""" # Read one_params file par = params.read('one_params') @@ -18,12 +18,40 @@ def _download_test_data(version=__version__): # Construct URL and call download url = '{}/{}/dlc_test_data_v{}.zip'.format(par.HTTP_DATA_SERVER, str(data_dir), '.'.join(version.split('.')[:-1])) - file_name = http_download_file(url, - cache_dir=local_path, - username=par.HTTP_DATA_SERVER_LOGIN, - password=par.HTTP_DATA_SERVER_PWD) + file_name = Path(http_download_file(url, + cache_dir=local_path, + username=par.HTTP_DATA_SERVER_LOGIN, + password=par.HTTP_DATA_SERVER_PWD)) # unzip file - shutil.unpack_archive(file_name, local_path) + test_dir = file_name.parent.joinpath(Path(file_name).stem) + if not test_dir.exists(): + shutil.unpack_archive(file_name, local_path) - return Path(file_name[:-4]) + return Path(test_dir) + + +def _download_me_test_data(): + """Download test data from FlatIron.""" + # eid: cde63527-7f5a-4cc3-8ac2-215d82e7da26 + # Read one_params file + par = params.read('one_params') + data_dir = Path('integration', 'dlc', 'test_data') + + # Create target directory if it doesn't exist + local_path = Path(par.CACHE_DIR).joinpath(data_dir) + local_path.mkdir(exist_ok=True, parents=True) + + # Construct URL and call download + url = '{}/{}/me_test_data.zip'.format(par.HTTP_DATA_SERVER, str(data_dir)) + file_name = Path(http_download_file(url, + cache_dir=local_path, + username=par.HTTP_DATA_SERVER_LOGIN, + password=par.HTTP_DATA_SERVER_PWD)) + + # unzip file + test_dir = file_name.parent.joinpath(Path(file_name).stem) + if not test_dir.exists(): + shutil.unpack_archive(file_name, local_path) + + return Path(test_dir) diff --git a/iblvideo/tests/test_choiceworld.py b/iblvideo/tests/test_choiceworld.py index 41e7570..3f6b6da 100644 --- a/iblvideo/tests/test_choiceworld.py +++ b/iblvideo/tests/test_choiceworld.py @@ -3,20 +3,20 @@ import pandas as pd from iblvideo.choiceworld import dlc from iblvideo.weights import download_weights -from iblvideo.tests.download_test_data import _download_test_data +from iblvideo.tests import _download_dlc_test_data from iblvideo import __version__ def test_dlc(version=__version__): - test_data = _download_test_data() + test_data = _download_dlc_test_data() path_dlc = download_weights(version=version) for cam in ['body', 'left', 'right']: file_mp4 = test_data.joinpath('input', f'_iblrig_{cam}Camera.raw.mp4') tmp_dir = test_data.joinpath('input', f'dlc_tmp_iblrig_{cam}Camera.raw') - out_file = dlc(file_mp4, path_dlc) + out_file, _ = dlc(file_mp4, path_dlc) assert out_file assert (tmp_dir.is_dir() is False) diff --git a/iblvideo/tests/test_motion_energy.py b/iblvideo/tests/test_motion_energy.py new file mode 100644 index 0000000..2aa044a --- /dev/null +++ b/iblvideo/tests/test_motion_energy.py @@ -0,0 +1,31 @@ +import os +import numpy as np +from iblvideo.motion_energy import motion_energy +from iblvideo.tests import _download_me_test_data + + +def test_motion_energy(): + + test_data = _download_me_test_data() + for cam in ['body', 'left', 'right']: + print(f"Running test for {cam}") + ctrl_me = np.load(test_data.joinpath(f'output/{cam}Camera.ROIMotionEnergy.npy')) + ctrl_roi = np.load(test_data.joinpath(f'output/{cam}ROIMotionEnergy.position.npy')) + dlc_pqt = test_data.joinpath(f'alf/_ibl_{cam}Camera.dlc.pqt') + + # Test with all frames + me_file, roi_file = motion_energy(test_data, dlc_pqt, frames=None) + test_me = np.load(me_file) + test_roi = np.load(roi_file) + assert all(test_me == ctrl_me) + assert all(test_roi == ctrl_roi) + + # Test with frame chunking + me_file, roi_file = motion_energy(test_data, dlc_pqt, frames=70) + test_me = np.load(me_file) + test_roi = np.load(roi_file) + assert all(test_me == ctrl_me) + assert all(test_roi == ctrl_roi) + + os.remove(me_file) + os.remove(roi_file)