Bridging the poses of humans and tetrapods with one general model.
- Generalized model: Unipose is a generalized model that can be applied to any tetrapod species.
- Joint datasets made easy: Unified datasets for humans and animals are hard to come by. Unipose provides a unified dataset class for humans and animals.
- Code quality: Unipose utilizes a modern python development toolchain using poetry for packaging, black for code style enforcement, pytest for automated testing and pydoctor for documentation.
Python 3.9 or later is required since we use plan to use dataclasses in the future and we heavily rely on type hints as updated in PEP 585 and PEP 586 (Literal
type).
During development, we use PyTorch 1.13.0 and torchvision 0.14.0.
The package is currently not published to PyPI. To install, clone the repository and run pip install -e .
in the root directory:
git clone git@github.com:Gennadiyev/unipose.git
cd unipose
pip install -e .
Or if you use poetry:
git clone git@github.com:Gennadiyev/unipose.git
cd unipose
poetry add -D .
Or if you use conda:
git clone git@github.com:Gennadiyev/unipose.git
cd unipose
conda env create -f environment.yaml -n unipose
conda activate unipose
Our models are not that heavy (~169 MB each). Fetch one from the table below and put it somewhere ($MODEL_PATH
).
Model | Specialty |
---|---|
model_run-5dd8_ep-60.pth |
For humans and tetrapods. Known to perform poorly if the pose is partially occluded or looks strange. |
model_run-cd37_ep-100.pth |
For humans and tetrapods. More stable than model_run-5dd8_ep-60.pth but not as accurate in confidence value. |
Our test server is out! It is a flask server that can be used to test the models. Especially useful during training since it scans all models under the model directory and let users pick any model they want. Not optimized for multi-user cases, if more than 1 user is perfoming inference, the server will queue them (expectedly).
-
Edit the
server/config.json
file to change the font path, port and model path. Usually the directory containing$MODEL_PATH
. -
Run
python3 server/server.py server/config.json
.
Note: Do not change the naming scheme of the model file. The server will automatically detect the model epoch and run ID.
You can also visualize the results with a featureful visualization script scripts/visualize.py
.
Example with model_run-5dd8_ep-60.pth
, using an image sampled from COCO dataset:
python scripts/visualize.py --dataset coco --checkpoint models/model_run-5dd8_ep-60.pth --dataset_path datasets/coco
Example using the same checkpoint with your own image:
python scripts/visualize.py --checkpoint models/model_run-5dd8_ep-60.pth --image_path test_images/your_own_image.jpeg
Use python3 scripts/visualize.py --help
to see the full usage.
Hidden feature! The
image_path
also supports a directory path. The script will visualize all images in the directory, and output a GIF animation! Files under the directory must follow the naming scheme*-{id}.*
or*_{id}.*
(more specifically in regex,^.+[-_]\d+\..+$
), or the output GIF may have unsorted frames.
Most datasets we use have COCO-style annotations. We use the unipose.cocotools
module to load the annotations.
coco
├── annotations/
│ ├── captions_train2017.json
│ ├── captions_val2017.json
│ ├── instances_train2017.json
│ ├── instances_val2017.json
│ ├── person_keypoints_train2017.json
│ └── person_keypoints_val2017.json
├── train2017/
│ ├── 000000000009.jpg
│ └── ...
└── val2017/
├── 000000000139.jpg
└── ...
(More datasets are coming soon.)
Our model is configuration-free and is trained with a single command:
python scripts/train.py --coco --output_dir exp
The log, model checkpoints, optimizer states will be saved to the exp
directory. Use python3 scripts/train.py --help
to see the full usage.
Since more than one dataset can be selected for one training session, the
--coco
flag can be used together with other dataset flags. Seescripts/train.py
for more details.
Currrently multi-GPU training is not supported. We are working on it.
Contributions are welcome! Please see CONTRIBUTING.md for more information.
There will be follow-up papers and code for a similar purpose. Stay tuned!
The developers use a few local mirrors to speed up conda package downloading. If you are not using the mirrors, you can remove the mirror URLs from the channels
section in environment.yaml
and use the default channels. Do not remove the conda-forge
channel since it is required for pytorch
and torchvision
.
In the latest version of UniPose, we added unipose.cocotools
as an adapted version of pycocotools
that supports our use case. However, the original pycocotools
is still required. Please update to the latest version of UniPose if you are using an older version.
This project has the same name as bmartacho/UniPose. However, this project is not affiliated with the project or paper in any way.