Track Stitching

README.md

@@ -2,7 +2,8 @@
 
 [![CI](https://github.com/talmolab/dreem/actions/workflows/ci.yml/badge.svg)](https://github.com/talmolab/dreem/actions/workflows/ci.yml)
 [![codecov](https://codecov.io/gh/talmolab/dreem/branch/main/graph/badge.svg?token=Sj8kIFl3pi)](https://codecov.io/gh/talmolab/dreem)
-<!-- [![Release](https://img.shields.io/github/v/release/talmolab/dreem?label=Latest)](https://github.com/talmolab/dream/releases/)
+[![Documentation](https://img.shields.io/badge/Documentation-dreem.sleap.ai-lightgrey)](https://dreem.sleap.ai)
+<!-- [![Release](https://img.shields.io/github/v/release/talmolab/dreem?label=Latest)](https://github.com/talmolab/dreem/releases/)
 [![PyPI](https://img.shields.io/pypi/v/dreem?label=PyPI)](https://pypi.org/project/dreem)
 ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/dreem) -->
 
@@ -148,13 +149,13 @@ The input into our training script is a `.yaml` file that contains all the param
 Once you have your config file and dataset set up, training is as easy as running
 
 ```bash
-python /path/to/dreem/training/train.py --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM]
+dreem-train --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM]
 ```
 where `CONFIG_DIR` is the directory that `hydra` should search for the `config.yaml` and `CONFIG_STEM` is the name of the config without the `.yaml` extension.
 
 e.g. If I have a config file called `base.yaml` inside my `/home/aaprasad/dreem_configs` directory I can call
 ```bash
-python /home/aaprasad/dreem/training/train.py --config-base=/home/aaprasad/dreem_configs --config-name=base
+dreem-train --config-base=/home/aaprasad/dreem_configs --config-name=base
 ```
 
 > Note: you can use relative paths as well but may be a bit riskier so we recommend absolute paths whenever possible.
@@ -168,26 +169,26 @@ Instead of changing the `base.yaml` file every time you want to run a different
 For overriding specific params with a sub-config, you can specify a `params_config` key and path in your `config.yaml` or run
 
 ```bash
-python /path/to/dreem/training/train.py --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] ++params_config="/path/to/params.yaml"
+dreem-train --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] ++params_config="/path/to/params.yaml"
 ```
 
 e.g. If I have a `params_to_override.yaml` file inside my `/home/aaprasad/dreem_configs` directory that contains a only a small selection of parameters that I'd like to override, I can run:
 
 ```bash
-python /home/aaprasad/dreem/training/train.py --config-base=/home/aaprasad/dreem_configs --config-name=base ++params_config=/home/aaprasad/dreem_configs/params_to_override.yaml
+dreem-train --config-base=/home/aaprasad/dreem_configs --config-name=base ++params_config=/home/aaprasad/dreem_configs/params_to_override.yaml
 ```
 
 ##### CLI-based override
 For directly overriding a specific param via the command line directly you can use the `section.param=key` syntax as so:
 
 ```bash
-python /path/to/dreem/training/train.py --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] section.param=value
+dreem-train --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] section.param=value
 ```
 
 e.g If now I want to override a couple parameters again, say change the number of attention heads and change the name of this run in my logger, I can pass `model.head=3` and `logger.name="test_nheads=3"` into 
 
 ```bash
-python /home/aaprasad/dreem/training/train.py --config-base=/home/aaprasad/dreem_configs --config-name=base model.nhead=3 logger.name="test_nheads=3"
+dreem-train --config-base=/home/aaprasad/dreem_configs --config-name=base model.nhead=3 logger.name="test_nheads=3"
 ```
 > Note: using the `section.param` syntax for CLI override will only override if the parameter exists in your config file otherwise an error will be thrown
 > if you'd like add a new parameter you can add `++` to the front of `section.param` e.g `++model.nhead=3`, However in this case, if the parameter exists in the config it will throw an error.
@@ -199,11 +200,11 @@ See [here](https://hydra.cc/docs/advanced/override_grammar/basic/) for more info
 
 > Note: You will usually only need the direct override or the file-based override however you can technically do both via 
 > ```bash
-> python /path/to/dreem/training/train.py --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] ++params_config="/path/to/params.yaml" section.param=value
+> dreem-train --config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM] ++params_config="/path/to/params.yaml" section.param=value
 > ```
 > e.g. alongside the `params_to_override.yaml`
 > ```bash
-> python /home/aaprasad/dreem/training/train.py --config-base=/home/aaprasad/dreem_configs --config-name=base ++params_config=/home/aaprasad/dreem_configs/params_to_override.yaml model.nhead=3 logger.name="test_nheads=3"
+> dreem-train --config-base=/home/aaprasad/dreem_configs --config-name=base ++params_config=/home/aaprasad/dreem_configs/params_to_override.yaml model.nhead=3 logger.name="test_nheads=3"
 > ```
 > but the `--config-base=[CONFIG_DIR] --config-name=[BASE_CONFIG_STEM]` arguments are always required. 
 > However, be careful to ensure there are no conflicts when combining CLI and file-based override syntax
@@ -260,13 +261,13 @@ Please see the [README](dreem/inference/configs/inference.yaml) in `dreem/infere
 Just like training we can use the hydra syntax for specifying arguments via the cli. Thus you can run inference via:
 
 ```bash
-python /path/to/dreem/inference/inference.py --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM]
+dreem-track --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM]
 ```
 
 e.g. If I had an inference config called `track.yaml` inside `/home/aaprasad/dreem_configs` 
 
 ```bash
-python /home/aaprasad/dreem/inference/inference.py --config-base=/home/aaprasad/dreem_configs --config-name=track
+dreem-track --config-base=/home/aaprasad/dreem_configs --config-name=track
 ```
 
 ##### Overriding Parameters.
@@ -275,11 +276,11 @@ Because there aren't as many parameters during inference as during training we r
 In order to override params via the CLI, we can use the same `hydra` `section.param` syntax:
 
 ```bash
-python /path/to/dreem/inference/inference.py --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM] section.param=[VALUE]
+dreem-track --config-base=[CONFIG_DIR] --config-name=[CONFIG_STEM] section.param=[VALUE]
 ```
 e.g if I want to set the window size of the tracker to 32 instead of 8 through `tracker.window_size=32` and use a different model saved in `/home/aaprasad/models/new_best.ckpt` I can do:
 ```bash
-python /home/aaprasad/dreem/inference/inference.py --config-base=/home/aaprasad/dreem_configs --config-name=track ckpt_path="/home/aaprasad/models/new_best.ckpt" tracker.window_size=32`
+dreem-track --config-base=/home/aaprasad/dreem_configs --config-name=track ckpt_path="/home/aaprasad/models/new_best.ckpt" tracker.window_size=32`
 ```
 #### Output
 This will run inference on the videos/detections you specified in the `dataset.test_dataset` section of the config and save the tracks to individual `[VID_NAME].dreem_inference.slp` files. If an `outdir` is specified in the config it will save to  `./[OUTDIR]/[VID_NAME].dreem_inference.slp`, otherwise it will just save to `./results/[VID_NAME].dreem_inference.slp`. Now you can load the file with `sleap-io` and do what you please!

docs/cli.md

@@ -0,0 +1,104 @@
+# Command-line Interface
+
+DREEM provides several types of functionality accessible through a command prompt.
+
+## Training
+
+### `dreem-train`
+
+`dreem-train` is the command-line interface for [training](./reference/dreem/training/train.md). Use this for training on a remote machine/cluster/colab notebook instead of using the API directly.
+
+#### Usage
+
+```
+usage: dreem-train [-h] [--hydra-help] [--config-dir] [--config-name] [+params_config] [+batch_config]
+
+positional arguments:
+    --config-dir    Path to configs dir
+    --config-name   Name of the .yaml config file stored in config-dir without the .yaml extension
+
+optional arguments:
+    -h, --help      Shows the application's help and exit.
+    --hydra-help    Shows Hydra specific flags (recommended over -h) 
+    +params_config  Path to .yaml file containing subset of params to override
+    +batch_config   Path to .csv file where each row indicates params to override for a single task in a batch job
+```
+
+See [the usage guide](./usage.md#step-21-setup-config) for a more in-depth explanation on how to use `dreem-train` and see [the training config walkthrough](./configs/training.md) for all available parameters.
+
+### `dreem-track`
+
+`dreem-track` is the command-line interface for [inference](./reference/dreem/inference/track.md). Use this for tracking using a pretrained model on a remote machine/cluster/notebook instead of using the API directly.
+
+#### Usage
+
+```
+usage: dreem-track [-h] [--hydra-help] [--config-dir] [--config-name] [+params_config] [+batch_config]
+
+positional arguments:
+    --config-dir    Path to configs dir
+    --config-name   Name of the .yaml config file stored in config-dir without the .yaml extension
+
+optional arguments:
+    -h, --help      Shows the application's help and exit.
+    --hydra-help    Shows Hydra specific flags (recommended over -h) 
+    +params_config  Path to .yaml file containing subset of params to override
+    +batch_config   Path to .csv file where each row indicates params to override for a single task in a batch job
+```
+See [the usage guide](./usage.md#step-2-set-up-config) for a more in-depth explanation on how to use `dreem-track` and see [the inference config walkthrough](./configs/inference.md) for all available parameters.
+
+### `dreem-visualize`
+
+`dreem-visualize` is the command-line utility for [generating annotated videos](./reference/dreem/io/visualize.md#dreem.io.visualize.annotate_video) based on DREEM output
+
+#### Usage
+
+```
+usage: dreem-visualize [-h] [--hydra-help] [+labels_path] [+vid_path] [+save_path] [+annotate.key] [--config-path] [+annotate.color_palette="tab20"] [+annotate.trails=2] [+annotate.boxes=64] [+annotate.names=True] [+annotate.track_scores=0.5] [+annotate.centroids=4] [+annotate.poses=False] [+annotate.fps=30] [+annotate.alpha=0.2]
+
+positional arguments:
+    +labels_path                The path to the dreem tracks as a `.csv` file containing the following keys:
+                                    - "X": the X position of the instance in pixels
+                                    - "Y": the Y position of the instance in pixels
+                                    - "Frame_id": The frame of the video in which the instance occurs
+                                    - "Pred_track_id": The track id output from DREEM
+                                    - "Track_score": The trajectory score output from DREEM
+                                    - (Optional) "Gt_track_id": the gt track id for the instance (mostly used for debugging)
+                                where each row is an instance. See `dreem.inference.track.export_trajectories` 
+    +vid_path                   The path to the video file in an `imageio`-accepted format (See: https://imageio.readthedocs.io/en/v2.4.1/formats.html)
+    +save_path                  The path to where you want to store the annotated video (e.g "./tracked_videos/dreem_output.mp4")
+    +annotate.key               The key where labels are stored in the dataframe - mostly used for choosing whether to annotate based on pred or gt labels
+
+optional arguments
+    -h, --help                  Shows the application's help and exit.
+    --hydra-help                Shows Hydra specific flags (recommended over -h) 
+    --config-path               The path to the config .yaml file containing params for annotating. We recommend using this file in lieu of any "annotate.*" arguments.
+    +annotate.color_palette     The matplotlib colorpalette to use for annotating the video. Defaults to `tab10`
+    +annotate.trails            The size of the track trails. If the trail size <=0 or none then it is not used.
+    +annotate.boxes             The size of the bbox. If bbox size <= 0 or None then it is not added
+    +annotate.names             Whether or not to annotate with name
+    +annotate.centroids         The size of the centroid. If centroid size <=0 or None then it is not added.
+    +annotate.poses             Whether or not to annotate with poses
+    +annotate.fps               The frames-per-second of the annotated video
+    +annotate.alpha             The opacity of the centroid annotations.
+```
+
+As noted above, instead of specifying each individual `+annotate.*` param we recommend setting up a visualize config (e.g `./configs/visualize.yaml`) which looks like:
+
+```YAML
+annotate:
+    color_palette: "tab10"
+    trails: 5
+    boxes: (64, 64)
+    names: true
+    centroids: 10
+    poses: false
+    fps: 30
+    alpha: 0.2
+```
+
+and then running:
+
+```bash
+dreem-visualize +labels_path="/path/to/labels.csv" +vid_path="/path/to/vid.{VID EXTENSION} +annotate.key="Pred_track_id" --config-path="/path/to/configs/visualize.yaml"
+```

docs/index.md

@@ -2,6 +2,7 @@
 
 [![CI](https://github.com/talmolab/dreem/actions/workflows/ci.yml/badge.svg)](https://github.com/talmolab/dreem/actions/workflows/ci.yml)
 [![codecov](https://codecov.io/gh/talmolab/dreem/branch/main/graph/badge.svg?token=Sj8kIFl3pi)](https://codecov.io/gh/talmolab/dreem)
+[![code](https://img.shields.io/github/stars/talmolab/dreem)](https://github.com/talmolab/dreem)
 <!-- [![Release](https://img.shields.io/github/v/release/talmolab/dreem?label=Latest)](https://github.com/talmolab/dreem/releases/)
 [![PyPI](https://img.shields.io/pypi/v/dreem?label=PyPI)](https://pypi.org/project/dreem)
 ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/dreem) -->