Left: confidence-colored pupil mask and estimated center — Right: evolving pupil diameter, center, and speed

See media/README.md to regenerate and review the demo GIF.

Pupil Analysis Pipeline

This package runs a full pipeline for mouse pupil segmentation and size estimation using a trained UNet model.
You can start directly from a video file or from an existing folder of extracted frames. To obtain expected results, the video or images provided should have at least the majority of the eye contained in the 148 x 148 pixel area in the center of the frames. This is crucial to getting good results because the model was trained on 148 x 148 centered eye images.

📦 Installation

It is recommended that you first create a dedicated virtual environment, for example with Miniconda:

conda create -n mouse_pupil_analysis python=3.12
conda activate mouse_pupil_analysis

Then install the package:

pip install mouse-pupil-analysis

The trained model checkpoint ships with the package, so there is nothing else to download.

Note on names. The repository and distribution are mouse-pupil-analysis, the Python import is mouse_pupil_analysis, and the established console commands remain run-pupil-analysis and extract-frames. The shorter name pupil-tracking on PyPI belongs to an unrelated project by a different author, and it installs its own pupil_tracking module, so this project deliberately claims no pupil_tracking import namespace of any kind. Existing Conda environments do not need to be renamed; environment names are local only.

GPU / CPU builds of PyTorch

On Windows and macOS, the command above installs a CPU-only build of PyTorch, which is all this package needs to run. No action required.

On Linux, the default PyPI wheel bundles CUDA and is several times larger. If you do not have an NVIDIA GPU, install the CPU-only build first to avoid the download:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip install mouse-pupil-analysis

To use an NVIDIA GPU, install a matching CUDA build first. This package requires torch>=2.8, which is served by the cu126, cu128, and cu129 indexes; older indexes such as cu124 stop at PyTorch 2.6 and will not satisfy that floor. Pick the index matching your driver with the official selector:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
pip install mouse-pupil-analysis

Inference selects the GPU automatically when one is available and falls back to CPU otherwise.

Installing for development

git clone https://github.com/yzhaoinuw/mouse-pupil-analysis.git
cd mouse-pupil-analysis
pip install -e ".[dev]"

🏃 Basic Usage

After installation, you can run pupil analysis on a video like so

run-pupil-analysis --video_path /path/to/movie.avi

This will:

1. Extract evenly spaced frames from the video into a folder like movie_frames/
2. Run pupil segmentation and diameter estimation on those frames
3. Save the results (CSV + plot) into movie_result/

To calculate pupil-center position and velocity from every encoded frame, add --calculate_velocity and provide the actual acquisition rate when it differs from the video playback rate:

run-pupil-analysis \
  --video_path data/mouse1.avi \
  --calculate_velocity \
  --acquisition_fps 33.3333333333

Velocity mode calculates timestamps from the original source-frame index and the acquisition rate. It does not assume that the FPS stored in the video container always represents experimental time.

For the end-to-end methodology—from segmentation probabilities through pupil-center quality control and velocity—see Segmentation-To-Velocity Method.

Try the Included Sample Data

The repository includes a compact set of real images and hand-labeled masks under sample_data/. It supports clone-and-run checks of uncropped-frame inference, overlays, paired training data, augmentation, and a 31-frame pupil-velocity sequence.

run-pupil-analysis \
  --image_dir sample_data/velocity_frames \
  --result_dir results/sample_velocity \
  --output_mask_dir results/sample_velocity/overlays \
  --calculate_velocity \
  --acquisition_fps 97

The fixture is intended for workflow exploration and debugging, not scientific model evaluation or useful model training. See the sample-data guide for the uncropped-frame and training examples.

⚙️ Key Arguments

💡 Examples

From a video (auto frame extraction):

run-pupil-analysis --video_path data/mouse1.avi

From an existing folder of frames:

run-pupil-analysis --image_dir data/mouse1_frames

With custom output locations and segmentation masks:

run-pupil-analysis \
  --video_path data/mouse1.avi \
  --out_dir data/frames_mouse1 \
  --result_dir data/results_mouse1 \
  --output_mask_dir data/masks_mouse1

🐍 Python API

Everything the CLI does is available from Python, which is usually more convenient inside a notebook or a larger analysis script. The results come back as a DataFrame, so there is no need to read the CSV back in.

from mouse_pupil_analysis import analyze_video

result = analyze_video("data/mouse1.avi")
print(result.analysis_table.head())
print(result.csv_path, result.plot_path)

Velocity mode and every CLI flag are keyword arguments:

result = analyze_video(
    "data/mouse1.avi",
    calculate_velocity=True,
    acquisition_fps=33.3333333333,
    output_mask_dir="data/masks_mouse1",
)

usable = result.analysis_table.query("tracking_status != 'invalid'")
print(f"{len(usable)} of {len(result.analysis_table)} frames usable")

To start from frames you already extracted, use analyze_frames instead:

from mouse_pupil_analysis import analyze_frames

result = analyze_frames(
    "data/mouse1_frames",
    calculate_velocity=True,
    acquisition_fps=33.3333333333,
)

result is an AnalysisResult with these fields:

For repeated runs with shared settings, build an AnalysisConfig once and pass it to run_analysis:

from mouse_pupil_analysis import AnalysisConfig, run_analysis

for video in Path("data").glob("*.avi"):
    run_analysis(AnalysisConfig(video_path=video, pred_thresh=0.75))

Library code logs rather than prints, so it stays quiet by default. To see the same progress messages the CLI shows:

import logging
logging.basicConfig(level=logging.INFO)

📦 Output Files

After running, you’ll typically find:

Units

Pupil diameter is reported twice, in two different units:

• estimated_pupil_diameter is measured in the 148 x 148 model image. Because every frame is rescaled to that size, this value is not comparable between recordings with different resolution or cropping. It is kept for continuity with earlier results.
• pupil_diameter_input_pixels inverts the resize-and-pad geometry to express the same measurement at the scale of the image you supplied. With --video_path that is the source video frame. With --image_dir it is whatever you prepared: if your frames were already cropped or resized to 148 x 148, this column equals estimated_pupil_diameter.

Both are equivalent-circle diameters: the diameter of a circle whose area matches the segmented pupil mask, sqrt(4 / pi * area).

Neither column is calibrated. pupil_diameter_input_pixels removes the model's rescaling, but two recordings still only compare directly if their optics and working distance match. Otherwise apply your own per-recording scale factor.

Pupil-center coordinates and speed use the same input-image pixel scale as pupil_diameter_input_pixels: the source video frame with --video_path, and whatever you supplied with --image_dir. The x coordinate increases to the right and the y coordinate increases downward, and speed is in input-image pixels per second. The same calibration caveat applies, so speeds are only directly comparable between recordings whose optics and working distance match.

Neither unit is physical. Converting to millimeters requires a scale factor from your own optics, which this package does not attempt to infer.

The compact analysis CSV reports tracking_status as valid, warning, or invalid, with quality_reason identifying suspicious or rejected frames. Published center and speed fields are left empty when segmentation is rejected. Speed is also left empty when either adjacent frame is invalid or when source frames are not consecutive; the pipeline does not interpolate across these gaps. A warning remains usable, such as extra foreground components when the selected pupil component is still acceptable.

🧩 Typical Folder Structure

movie.avi
movie_frames/
    movie_00001.png
    movie_00002.png
    ...
movie_result/
    movie_pupil_analysis.csv
    movie_pupil_analysis.png

Citation

If you use this software in a paper or other scholarly work, please cite the version you used. GitHub renders citation metadata from CITATION.cff, and the “Cite this repository” button produces BibTeX directly.

Recommended citation:

Yue Zhao. mouse-pupil-analysis: Automated mouse pupil segmentation, diameter, and pupil-center velocity analysis using UNet. Version 0.2.0. https://github.com/yzhaoinuw/mouse-pupil-analysis

Zenodo archiving is enabled, but no DOI exists until the first GitHub Release is published; see RELEASING.md. Once Zenodo processes that release, cite the DOI of the specific version you ran so your analysis stays reproducible against that exact code. Until then, cite the version number and commit. See CHANGELOG.md for what changed between versions.

License

This project is released under the MIT License. See LICENSE.

Developer Notes

Model Training

The complete data-preparation, augmentation, fresh-training, fine-tuning, and checkpoint-promotion workflow is documented in training/README.md.

Making Training Data

Create two folders in mouse-pupil-analysis/, images_train/ and masks_train/ if you haven't. Place your training images in images_train/. Once you have done this once, you can just add new training images to images_train/.

1. In Terminal/Anaconda Powershell Prompt, activate your project environment, then run labelme.exe to open the labelme interface to label images.
2. After you are done, labelme should have saved your labels as json files in images_train/ along with your training images. Now run python .\training\labelme_json2png.py, which will create the masks (png files) and move them to masks_train/.
3. To create the validation set, create images_validation/ and masks_validation/ and then follow the same steps above, but remember to change dataset_type in training/labelme_json2png.py accordingly.
4. To start training the model, run python .\training\run_train.py. You can modify the hyperparameters in training/run_train.py as needed.