[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/talmolab/dreem/blob/main/examples/dreem-demo.ipynb)

# End-to-end demo

This notebook will walk you through the DREEM pipeline end to end, from obtaining data to training a model, evaluating on a held-out dataset, and visualizing the results. Here, we use the **CLI** for training, tracking, and evaluation; an **Advanced Usage** section shows how to use configuration files for finer control.

To run this demo, we have provided sample data and configurations. The data used in this demo is small enough to be run on a CPU.

#### Install DREEM

```python
!uv pip install dreem-track
```

```python
import pandas as pd
import numpy as np
import os
import glob
import yaml
import sleap_io as sio
```

## Download data

```python
!hf download talmolab/sample-flies --repo-type dataset --local-dir ./data
```

## Training

We just need to specify the data paths and the crop size to train the model. The crop size is the size (in pixels) of the square box we make around the instance. It should be of the order of the size of the instance.

Lets figure out an appropriate crop size. We can do this by loading a frame and experimenting with different sizes. For this data, a crop size of 70 seems appropriate.

```python
crop_size = 70  # Adjust this to see what the bounding box looks like
```

```python
slp_files = glob.glob("./data/train/*.slp")
video_files = glob.glob("./data/train/*.mp4")
labels = sio.load_slp(slp_files[0])
video = imageio.get_reader(video_files[0])
frame = video.get_data(0)
centers = []
for instance in labels[0].instances:
    centers.append(np.nanmean(instance.numpy(), axis=0))
centers = np.array(centers)
plt.imshow(frame)
plt.scatter(centers[:, 0], centers[:, 1])
# Draw a square box centered at each centroid with a given crop size (e.g., 70)
for cx, cy in centers:
    top_left = (int(cx - crop_size // 2), int(cy - crop_size // 2))
    rect = plt.Rectangle(top_left, crop_size, crop_size, linewidth=2, edgecolor='red', facecolor='none')
    plt.gca().add_patch(rect)
plt.show()
```

![Crop size preview: frame with centroids and bounding boxes](crop-size.png)

### That's it! Now we can train the model

```python
!dreem train ./data/train --val-dir ./data/val --crop-size {crop_size}
```

## Tracking
Here we run tracking on a video with **no** ground truth labels using the model we just trained. Note that we're using a command line argument to set the maximum number of tracks to 2 since its a 2 flies dataset. You can run the help command below to see all the options.

```python
!dreem track --help
```

```python
models_dir = "./models/dreem_train"
ckpt_files = glob.glob(os.path.join(models_dir, "*final*.ckpt"))
final_ckpt = ckpt_files[0]  # Pick the first one found
print(f"Using checkpoint: {final_ckpt}")
```

```python
!dreem track ./data/inference --checkpoint {final_ckpt} --output ./results --crop-size {crop_size} --max-tracks 2
```

## Evaluate the tracking results
#### Here we run inference on a video **with** ground truth labels. This enables us to compute metrics for our tracking results.

```python
!dreem eval ./data/test --checkpoint {final_ckpt} --output "./eval-results" --crop-size {crop_size} --max-tracks 2
```

## And we're done!
You can take a look at the tracking metrics in the motmetrics.csv file in the directory you chose to save the results to.

## Advanced Usage
Some parameters cannot be set through CLI overrides, particularly for training. This section demonstrates the use of configuration files to achieve fine grained control over all training parameters that Pytorch Lightning offers.

#### Setup configuration file
Here we override the default training parameters to use different augmentations, minimum epochs, and optimizer parameters. You can edit these directly in the config file at the path below. Try out some of your own augmentations! We use [Albumentations](https://albumentations.ai/) for augmentations.

```python
train_config_path = "./data/configs/base.yaml"
with open(train_config_path, "r") as f:
    config_yaml = yaml.safe_load(f)
print(yaml.dump(config_yaml, default_flow_style=False))
```

### Train the model with configs

```python
!dreem train ./data/train --val-dir ./data/val --crop-size {crop_size} --config ./data/configs/base.yaml
```

### Run inference with configs

Here we use inference overrides for more control over the tracking process. Specifically, we use max_center_dist, which limits how far each instance is allowed to be frame over frame in order to be considered the same instance, and a confidence threshold to flag low confidence predictions for manual review.

```python
track_config_path = "./data/configs/inference.yaml"
with open(track_config_path, "r") as f:
    config_yaml = yaml.safe_load(f)
print(yaml.dump(config_yaml, default_flow_style=False))
```

```python
!dreem track ./data/inference --checkpoint {final_ckpt} --output ./results --crop-size {crop_size} --config ./data/configs/inference.yaml
```

## Visualize the results
**Option A – DREEM Visualizer (browser-based)**  
Download your tracking results, and head to the [live visualizer](visualizer.md) to visualize your tracking results in your browser without any data leaving your machine.

**Option B – SLEAP GUI (full pose keypoints)**  
Install SLEAP ([https://docs.sleap.ai/latest/](https://docs.sleap.ai/latest/)) and open the output `.slp` in SLEAP:
```bash
sleap-label results/<your_output_file>.slp
```
The SLEAP GUI may not render on a remote server.