Towards Monopose

.github/workflows/pytest.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9"]
+        python-version: ["3.10"]
 
     steps:
       - uses: szenius/set-timezone@v1.1

.gitignore

@@ -14,4 +14,6 @@ test_images
 # apidocs generated by pydoctor
 apidocs
 .gittoken
-server_cache
+server_cache
+.conda
+.vscode

README.md

@@ -1,84 +1,73 @@
-# Unipose<a href="https://hexops.com"><img align="left" alt="Hexops logo" src="docs/unipose.svg" width="30%" style="padding-right:16px"></img></a>
+# Unipose<a href="https://github.com/Gennadiyev/unipose"><img align="left" alt="Unipose logo" src="docs/unipose.svg" width="30%" style="padding-right:16px"></img></a>
 
-[![stable-docs](https://shields.io/badge/docs-stable-blue.svg)](https://gennadiyev.github.io/unipose/apidocs) [![style-black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://black.vercel.app/) [![python>=3.9](https://img.shields.io/badge/python->=3.8-green.svg)](https://www.python.org/downloads/) [![license-MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://mit-license.org/) ![workflow-status](https://img.shields.io/github/actions/workflow/status/Gennadiyev/unipose/pytest.yml)
+[![stable-docs](https://shields.io/badge/docs-stable-blue.svg)](https://gennadiyev.github.io/unipose/apidocs) [![style-black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://black.vercel.app/) [![python>=3.10](https://img.shields.io/badge/python->=3.10-green.svg)](https://www.python.org/downloads/) [![license-MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://mit-license.org/) ![workflow-status](https://img.shields.io/github/actions/workflow/status/Gennadiyev/unipose/pytest.yml)
 
 **Bridging the poses of humans and tetrapods with one general model.**
 
 <!-- ## Capabilities
 
 - **Pose estimation**: Unipose can estimate the pose of any tetrapod with high accuracy.
 - **Pose tracking**: Unipose can track the pose of any tetrapod in a video sequence. (TODO)
-- **Animal Classification**: Unipose can classify the species of any tetrapod based on its pose. (TODO) -->
-
-<!-- ![Skeleton annotation of a giraffe](docs/skel_giraffe.png) ![Skeleton annotation of a human](docs/skel_human.png) -->
+- **Animal Classification**: Unipose can classify the species of any tetrapod based on its pose. (TODO) 
+-->
 
 <img src="docs/skel_giraffe.png" alt="Skeleton annotation of a giraffe" width="33%"></img> <img src="docs/skel_human.png" alt="Skeleton annotation of a human" width="33%"></img>
 
 ## Features
 
 - **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.
+- **Joint datasets made easy**: Unified datasets for humans and animals are hard to come by. We provide a unified dataset base class for humans and animals as well as literally anything.
 - **Code quality**: Unipose utilizes a modern python development toolchain using [poetry](https://python-poetry.org/) for packaging, [black](https://black.vercel.app/) for code style enforcement, [pytest](https://pytest.org/) for automated testing and [pydoctor](https://pydoctor.readthedocs.io/en/latest/) for documentation.
 
 ## Installation
 
 ### Prerequisites
 
-**Python 3.9** or later is required since we use plan to use [dataclasses](https://docs.python.org/3/library/dataclasses.html) in the future and we heavily rely on [type hints](https://docs.python.org/3/library/typing.html) as updated in [PEP 585](https://www.python.org/dev/peps/pep-0585/) and [PEP 586](https://peps.python.org/pep-0586/) (`Literal` type).
+**Python 3.10** is required since we are using `typing.TypeAlias` and `typing.Protocol` in the codebase. If you are wondering, we are also using [dataclasses](https://docs.python.org/3/library/dataclasses.html) and heavily rely on [type hints](https://docs.python.org/3/library/typing.html) as updated in [PEP 585](https://www.python.org/dev/peps/pep-0585/) and [PEP 586](https://peps.python.org/pep-0586/) (`Literal` type). 
 
 During development, we use PyTorch 1.13.0 and torchvision 0.14.0.
 
-### Install with pip
+### Install with conda
 
-The package is currently not published to PyPI. To install, clone the repository and run `pip install -e .` in the root directory:
+Or if you use [conda](https://docs.conda.io/en/latest/):
 
 ```bash
 git clone git@github.com:Gennadiyev/unipose.git
 cd unipose
-pip install -e .
+conda env create -f environment.yaml -n unipose
+conda activate unipose
 ```
 
-### Install with poetry
+### Install with pip
 
-Or if you use [poetry](https://python-poetry.org/):
+The package is currently not published to PyPI. To install, clone the repository and run `pip install -e .` in the root directory:
 
 ```bash
 git clone git@github.com:Gennadiyev/unipose.git
 cd unipose
-poetry add -D .
+pip install -e .
 ```
 
-### Install with conda
+### Install with poetry
 
-Or if you use [conda](https://docs.conda.io/en/latest/):
+Or if you use [poetry](https://python-poetry.org/):
 
 ```bash
 git clone git@github.com:Gennadiyev/unipose.git
 cd unipose
-conda env create -f environment.yaml -n unipose
-conda activate unipose
+poetry add -D .
 ```
 
-## Usage for testing
+## Usage
 
 ### Download pretrained models
 
-Our models are not that heavy (~169 MB each). Fetch one from the table below and put it somewhere (`$MODEL_PATH`).
+Our models are not that heavy (~169 MB each). Fetch one from [this Google Drive link](https://drive.google.com/drive/folders/1eJ9RyLHcezrxE02uHBNmrI0Wxg4OPkfQ). We suggest putting them under `$UNIPOSE_ROOT/models`.
 
 | Model | Specialty |
 | --- | --- |
-| [`model_run-5dd8_ep-60.pth`](https://drive.google.com/drive/folders/1eJ9RyLHcezrxE02uHBNmrI0Wxg4OPkfQ?usp=sharing) | For humans and tetrapods. Known to perform poorly if the pose is partially occluded or looks strange. |
-| [`model_run-cd37_ep-100.pth`](https://drive.google.com/file/d/1f3uAjAzKh5O4Gh6vNuDZ99e78JlhoBnT/view?usp=sharing) | For humans and tetrapods. More stable than `model_run-5dd8_ep-60.pth` but not as accurate in confidence value. |
-
-### Run a flask server for testing
-
-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).
-
-1. Edit the `server/config.json` file to change the font path, port and model path. Usually the directory containing `$MODEL_PATH`.
-
-2. 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.
+| [`model_run-5dd8_ep-60.pth`](https://drive.google.com/file/d/12NJ5EFBeyQdf-dtMDrV5Kg82Gq9vUl7L/view) | For humans and tetrapods. Known to perform poorly if the pose is partially occluded or looks strange. |
+| [`model_run-cd37_ep-100.pth`](https://drive.google.com/file/d/1f3uAjAzKh5O4Gh6vNuDZ99e78JlhoBnT/view) | For humans and tetrapods. More stable than `model_run-5dd8_ep-60.pth` but not as accurate in confidence value. |
 
 ### Visualize skeleton with script
 
@@ -87,19 +76,34 @@ You can also visualize the results with a featureful visualization script `scrip
 Example with `model_run-5dd8_ep-60.pth`, using an image sampled from COCO dataset:
 
 ```bash
-python scripts/visualize.py --dataset coco --checkpoint models/model_run-5dd8_ep-60.pth --dataset_path datasets/coco
+python scripts/visualize.py \
+    --dataset coco \
+    --dataset_path datasets/coco \
+    --checkpoint models/model_run-5dd8_ep-60.pth
 ```
 
 Example using the same checkpoint with your own image:
 
 ```bash
-python scripts/visualize.py --checkpoint models/model_run-5dd8_ep-60.pth --image_path test_images/your_own_image.jpeg
+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.
 
+### *(Beta)* Run a flask server for testing
+
+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 case**, if more than 1 user is perfoming inference, the server will queue them (expectedly).
+
+1. Edit the `server/config.json` file to change the font path, port and model path. Usually the directory containing of the experiment output path as specified in `output_dir` parameter during a [train run](#start-training).
+
+2. 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.
+
 ## Training
 
 ### Prepare datasets
@@ -146,6 +150,7 @@ The log, model checkpoints, optimizer states will be saved to the `exp` director
 - [Yikun Ji (Kunologist)](https://github.com/Gennadiyev)
 - [Qi Fan (fanqiNO1)](https://github.com/fanqiNO1)
 - [Yihang Qiu (GihhArwtw)](https://github.com/GihhArwtw)
+- [Xiangyun Rao (xyrrrrrrrr)](https://github.com/xyrrrrrrrr)
 
 ## Contributing
 
@@ -161,8 +166,8 @@ The developers use a few local mirrors to speed up conda package downloading. If
 
 ### Dataset issues
 
-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.
+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.
 
 ### Since it has the same name, is this project affiliated with that UniPose?
 
-This project has the same name as [bmartacho/UniPose](https://github.com/bmartacho/UniPose). However, this project is not affiliated with the project or [paper](https://arxiv.org/abs/2001.08095) in any way.
+We prefer to be styled as `Unipose` instead of `UniPose`. Despite the similar name as [bmartacho/UniPose](https://github.com/bmartacho/UniPose), we are not affiliated in any way.

environment.yaml

@@ -10,11 +10,12 @@ channels:
   - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
   - defaults
 dependencies:
-  - python=3.8
-  - pytorch=1.13.0
-  - torchvision=0.14.0
-  - pip=22.3.1
+  - python=3.10
+  - pip
   - pip:
+    - torch==1.13.1
+    - torchvision>=0.14,<0.15
+    - torchinfo
     - opencv-python>=4.6.0.0,<4.7.0.0
     - tqdm
     - orjson
@@ -25,4 +26,5 @@ dependencies:
     - imageio
     - easydict
     - pycocotools
-    - requests
+    - requests
+    - pytest

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "unipose"
-version = "0.0.2"
+version = "0.0.3"
 description = "Unified Pose Estimation Framework"
 authors = [
     {name = "Kunologist", email = "jiyikun2002@gmail.com"}
 ]
 dependencies = [
     "torch>=1.13.0,<2.0.0",
-    "torchvision>=0.14.0,<1.0.0",
+    "torchvision>=0.14.0,<0.15.0",
     "tqdm",
     "opencv-python>=4.0.0,<5.0.0",
     "orjson",
@@ -18,10 +18,11 @@ dependencies = [
     "imageio",
     "easydict",
     "requests",
-    "pycocotools"
+    "pycocotools",
+    "torchinfo"
 ]
 keywords = ["pose-estimation", "pose", "framework", "skeleton", "animal-pose"]
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 readme = "README.md"
 license = {text = "MIT"}
 
@@ -49,7 +50,7 @@ includes = ["unipose"]
 
 [tool.black]
 line-length = 120
-target-version = ['py38', 'py39', 'py310', 'py311']
+target-version = ['py310', 'py311']
 
 [tool.isort]
 profile = "black"
@@ -58,7 +59,7 @@ profile = "black"
 minversion = "6.0"
 addopts = "-ra -m \"not slow and not contains_absolute_path\""
 testpaths = [
-    "tests/*"
+    "./tests/*"
 ]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",

scripts/train.py

@@ -100,6 +100,8 @@ def get_abs_path(dir_path: str, create_if_not_exists: bool = False):
     parser.add_argument("--mpii_path", type=str, default="datasets/mpii")
     parser.add_argument("--animal_kingdom", action="store_true", default=False)
     parser.add_argument("--animal_kingdom_path", type=str, default="datasets/animal_kingdom")
+    parser.add_argument("--ap10k", action="store_true", default=False)
+    parser.add_argument("--ap10k_path", type=str, default="datasets/ap10k")
     parser.add_argument("-a", "--all", action="store_true", default=False)
     # Output configuration
     parser.add_argument("-o", "--output_dir", type=str, default="exp")
@@ -126,7 +128,7 @@ def get_abs_path(dir_path: str, create_if_not_exists: bool = False):
     logger.info("Using device {} (GPU {})", device, args.gpu)
 
     # Create model
-    logger.debug("Creating model...")
+    logger.debug("Creating UniPose model...")
     model = UniPose(13, resnet_layers=[3, 8, 36, 3])
     model = model.to(device)
 
@@ -144,11 +146,12 @@ def get_abs_path(dir_path: str, create_if_not_exists: bool = False):
         "animal_kingdom": get_abs_path(args.animal_kingdom_path),
         "mpii": get_abs_path(args.mpii_path),
         "coco": get_abs_path(args.coco_path),
+        "ap10k": get_abs_path(args.ap10k_path),
     }
 
     datasets = []
     logger.info("Loading datasets...")
-    from unipose.datasets import AnimalKingdomDataset, MPIIDataset, COCODataset, ConcatJointDataset
+    from unipose.datasets import AnimalKingdomDataset, MPIIDataset, COCODataset, ConcatJointDataset, AP10KDataset
 
     if args.animal_kingdom or args.all:
         _path = dataset_path.get("animal_kingdom")
@@ -167,8 +170,17 @@ def get_abs_path(dir_path: str, create_if_not_exists: bool = False):
         logger.debug("Loading COCO dataset from {}...", _path)
         dataset_coco = COCODataset(path=_path)
         datasets.append(dataset_coco)
+    if args.ap10k or args.all:
+        _path = dataset_path.get("ap10k")
+        logger.debug("Loading AP10K dataset from {}...", _path)
+        dataset_ap10k_1 = AP10KDataset(path=_path, split="train", sub_split=1)
+        dataset_ap10k_2 = AP10KDataset(path=_path, split="train", sub_split=2)
+        dataset_ap10k_3 = AP10KDataset(path=_path, split="train", sub_split=3)
+        datasets.append(dataset_ap10k_1)
+        datasets.append(dataset_ap10k_2)
+        datasets.append(dataset_ap10k_3)
     if len(datasets) == 0:
-        logger.error("No dataset selected: --coco, --mpii, --animal_kingdom or -a / --all must be specified")
+        logger.error("No dataset selected: --coco, --mpii, --animal_kingdom, --ap10k or -a / --all must be specified")
         exit(1)
     elif len(datasets) >= 2:
         dataset = ConcatJointDataset(datasets)

tests/datasets.py

@@ -44,4 +44,4 @@ def test_ap10k():
     data = dataset[0]
     assert data["unipose_keypoints"].shape == torch.Size([13, 2])
     assert data["bounding_box"].shape == torch.Size([4])
-    assert data["image"].shape[0] == 3
+    assert data["image"].shape[0] == 3

tests/test_model_v1.py

@@ -0,0 +1,15 @@
+from unipose.models import UniPose
+from unipose.layers import SEResNeXt, DUC
+
+def test_model_structure():
+    from torchinfo import summary
+    model = UniPose(17, channels=64, groups=32, reduction=16, resnet_layers=[3, 8, 36, 3], duc_layers=[4, 2, 1])
+    summary(model, input_size=(1, 3, 256, 256), depth=10)
+
+def test_layers():
+    from torchinfo import summary
+    model_duc = DUC(64, 256, 2, 3)
+    summary(model_duc, input_size=(1, 64, 4, 4))
+    model_seresnext = SEResNeXt(channels=64, groups=32, reduction=16, layers=[2, 2, 2, 2])
+    summary(model_seresnext, input_size=(1, 3, 256, 256), depth=10)
+

unipose/cocotools/coco.py

@@ -137,9 +137,9 @@ def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None):
         """
         Get ann ids that satisfy given filter conditions. default skips that filter
         :param imgIds  (int array)     : get anns for given imgs
-               catIds  (int array)     : get anns for given cats
-               areaRng (float array)   : get anns for given area range (e.g. [0 inf])
-               iscrowd (boolean)       : get anns for given crowd label (False or True)
+        :param catIds  (int array)     : get anns for given cats
+        :param areaRng (float array)   : get anns for given area range (e.g. [0 inf])
+        :param iscrowd (boolean)       : get anns for given crowd label (False or True)
         :return: ids (int array)       : integer array of ann ids
         """
         imgIds = imgIds if _isArrayLike(imgIds) else [imgIds]
@@ -381,7 +381,7 @@ def download(self, tarDir = None, imgIds = [] ):
         '''
         Download COCO images from mscoco.org server.
         :param tarDir (str): COCO results directory name
-               imgIds (list): images to be downloaded
+        :param imgIds (list): images to be downloaded
         :return:
         '''
         if tarDir is None: