Zebrafish Attention-based Trajectory Analyzer

ZATA (Zebrafish Attention-based Trajectory Analyzer) is an end-to-end system for analyzing zebrafish behavior from trajectory data using attention-based neural networks. It processes sequential input features and outputs a continuous confidence score, indicating the possible targer behavior locations in time frames.

🧠 Overview

The model takes in preprocessed trajectory data with sliding windows and trains a neural network to detect patterns. It's especially suited for experiments like analyzing fish swimming behavior under different conditions.

Project Structure

.
├── data/                  
│   ├── data_labels/          # Ground truth label files (.xlsx)
│   ├── data_trajectories/    # Input trajectories from DeepLabCut (.h5)       
│   └── temperature_record/   # Experiment temperature logs
│
├── outputs/                  # Saved models and inference results
│
├── utils/                    # Core Model Logic
│   ├── net.py                # Neural Network architecture
│   ├── detector_model.py     # Training & Evaluation loop
│   ├── data_preprocessing.py # Data loading and windowing
│   └── generate.py           # Video clip extraction
│
├── inference.py              # Script to detect behaviors
└── train.py                  # Script to train a new model

Quick Start

This guide walks you through how to prepare data and use the model for inference. For model training process, you can check the documents How to Train ZATA Model

0. Analyze Animal Trajectories

Use DeepLabCut to track animal movements and export the trajectory recordings.

1. Organize Your Data

Place your files in the following directory structure:

• Place your trajectory records .h5 files in the data/data_trajectories folder.

2. Inference

Run the script using:

bash inference.sh

where the parameters can be modified:

CKPT_PATH="./outputs/saved_models/H100.pt"
TRAJECTORY_FILE_PATH= [FILL IN]   # path to the trajectory .h5 file
VIDEO_FILE_PATH= [FILL IN]        # path to the corresponding video file
OUTPUT_SAVE_PATH="./results"      # location to save the inference results