kaiyangzhou/osnet

This repo contains the pre-trained weights of OSNet, specialized for person recognition (i.e., person re-identification).

Related work:

• Omni-Scale Feature Learning for Person Re-Identification
• Learning Generalisable Omni-Scale Representations for Person Re-Identification

Get started

Install the Torchreid package.

# cd to your preferred directory and clone this repo
git clone https://github.com/KaiyangZhou/deep-person-reid.git

# create environment
cd deep-person-reid/
conda create --name torchreid python=3.7
conda activate torchreid

# install dependencies
# make sure `which python` and `which pip` point to the correct path
pip install -r requirements.txt

# install torch and torchvision (select the proper cuda version to suit your machine)
conda install pytorch torchvision cudatoolkit=9.0 -c pytorch

# install torchreid (don't need to re-build it if you modify the source code)
python setup.py develop

Download OSNet models.

from huggingface_hub import snapshot_download

# This will download the entire repo (containing all models) to cache_dir
snapshot_download(repo_id="kaiyangzhou/osnet", cache_dir="./")

Use Torchreid as a feature extractor in your project.

from torchreid.utils import FeatureExtractor

extractor = FeatureExtractor(
    model_name='osnet_x1_0',
    model_path='a/b/c/model.pth',
    device='cuda'
)

image_list = [
    'a/b/c/image001.jpg',
    'a/b/c/image002.jpg',
    'a/b/c/image003.jpg',
    'a/b/c/image004.jpg',
    'a/b/c/image005.jpg'
]

features = extractor(image_list)
print(features.shape) # output (5, 512)

Model list

Models pre-trained on ImageNet are named in the following format: osnet_x<scale>_imagenet.pth.

Available re-id models:

• osnet_ain_x1_0_msmt17_256x128_amsgrad_ep50_lr0.0015_coslr_b64_fb10_softmax_labsmth_flip_jitter.pth
• osnet_ibn_x1_0_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip_jitter.pth
• osnet_x1_0_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip_jitter.pth
• osnet_x0_75_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip_jitter.pth
• osnet_x0_5_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip_jitter.pth
• osnet_x0_25_msmt17_combineall_256x128_amsgrad_ep150_stp60_lr0.0015_b64_fb10_softmax_labelsmooth_flip_jitter.pth

ain and ibn models are more generalizable. Please refer to https://arxiv.org/abs/1910.06827.

osnet_x1_0, osnet_x0_75, osnet_x0_5, and osnet_x0_25 are OSNet models of different sizes. Please refer to https://arxiv.org/abs/1905.00953.