Repository: bupt-ai-cz/HHCL-ReID
Branch: main
Commit: 4986b46251de
Files: 41
Total size: 128.1 KB
Directory structure:
gitextract_d7xkgx8j/
├── README.md
├── examples/
│ ├── test.py
│ └── train.py
├── hhcl/
│ ├── __init__.py
│ ├── datasets/
│ │ ├── __init__.py
│ │ ├── celebreid.py
│ │ ├── dukemtmcreid.py
│ │ ├── market1501.py
│ │ ├── msmt17.py
│ │ └── personx.py
│ ├── evaluation_metrics/
│ │ ├── __init__.py
│ │ ├── classification.py
│ │ └── ranking.py
│ ├── evaluators.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── cm.py
│ │ ├── dsbn.py
│ │ ├── kmeans.py
│ │ ├── losses.py
│ │ ├── pooling.py
│ │ ├── resnet.py
│ │ ├── resnet_ibn.py
│ │ ├── resnet_ibn_a.py
│ │ └── triplet.py
│ ├── trainers.py
│ └── utils/
│ ├── __init__.py
│ ├── data/
│ │ ├── __init__.py
│ │ ├── base_dataset.py
│ │ ├── preprocessor.py
│ │ ├── sampler.py
│ │ └── transforms.py
│ ├── faiss_rerank.py
│ ├── faiss_utils.py
│ ├── logging.py
│ ├── meters.py
│ ├── osutils.py
│ ├── rerank.py
│ └── serialization.py
├── requirements.txt
├── run.sh
└── setup.py
================================================
FILE CONTENTS
================================================
================================================
FILE: README.md
================================================
# HHCL-ReID 
[](https://twitter.com/intent/tweet?text=Codes%20for%20Our%20Paper:%20"Hard-sample%20Guided%20Hybrid%20Contrast%20Learning%20for%20Unsupervised%20PersonRe-Identification"%20&url=https://github.com/bupt-ai-cz/HHCL-ReID) [](https://paperswithcode.com/sota/unsupervised-person-re-identification-on-5?p=hard-sample-guided-hybrid-contrast-learning) [](https://paperswithcode.com/sota/unsupervised-person-re-identification-on-4?p=hard-sample-guided-hybrid-contrast-learning)
This repository is the official implementation of our paper "[Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-Identification](https://arxiv.org/abs/2109.12333)!".
## Requirements
---
git clone https://github.com/bupt-ai-cz/HHCL-ReID.git
cd HHCL-ReID
pip install -r requirements.txt
python setup.py develop
## Prepare Datasets
---
Download the datasets Market-1501,MSMT17,DukeMTMC-reID from this [link](https://drive.google.com/file/d/19oWiYGjTgouFMK_psZvH8ysDGQ1KUbk-/view?usp=sharing) and unzip them under the directory like:
HHCL-ReID/examples/data
├── market1501
│ └── Market-1501-v15.09.15
└── dukemtmcreid
└── DukeMTMC-reID
Prepare ImageNet Pre-trained Models for IBN-Net
When training with the backbone of [IBN-ResNet](https://arxiv.org/abs/1807.09441), you need to download the ImageNet-pretrained model from this [link](https://drive.google.com/drive/folders/1thS2B8UOSBi_cJX6zRy6YYRwz_nVFI_S) and save it under the path of `examples/pretrained/`.
```
HHCL-ReID/examples
└── pretrained
└── resnet50_ibn_a.pth.tar
```
## Training
---
We utilize 4 GTX-2080TI GPUs for training. Examples:
Market-1501:
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d market1501 --iters 200 --eps 0.45 --momentum 0.1 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_avg_cmhybrid
DukeMTMC-reID:
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d dukemtmcreid --iters 200 --eps 0.6 --momentum 0.1 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_avg_cmhybrid
- use `-a resnet50` (default) for the backbone of ResNet-50, and `-a resnet_ibn50a` for the backbone of IBN-ResNet;
- use `--pooling-type gem` for Generalized Mean Pooling (GEM) pooling and `--smooth` for label smoothing.
## Evaluation
---
To evaluate my model on ImageNet, run:
CUDA_VISIBLE_DEVICES=0 python examples/test.py -d $DATASET --resume $PATH --pooling-type avg
## Results
---
Our model achieves the following performance on :
| Dataset | Market1501 | | | | DukeMTMC-reID | | | |
| ------------------ | ---------- | ---- | ---- | ---- | ------------- | ---- | ---- | ---- |
| Setting | mAP | R1 | R5 | R10 | mAP | R1 | R5 | R10 |
| Fully Unsupervised | 84.2 | 93.4 | 97.7 | 98.5 | 73.3 | 85.1 | 92.4 | 94.6 |
| Supervised | 87.2 | 94.6 | 98.5 | 99.1 | 80.0 | 89.8 | 95.2 | 96.7 |
You can download the above models in the paper from [Google Drive](https://drive.google.com/drive/folders/1WQw7wD2Mu_1SKl07_NdKvrYf2xrs3CEZ)
## Citation
---
If you find this code useful for your research, please cite our paper
```
@article{hu2021hard,
title={Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-Identification},
author={Hu, Zheng and Zhu, Chuang and He, Gang},
journal={arXiv preprint arXiv:2109.12333},
year={2021}
}
```
## Acknowledgements
---
This project is not possible without multiple great opensourced codebases. We list them below.
- [SpCL](https://github.com/yxgeee/SpCL)
- [cluster-contrast-reid](https://github.com/alibaba/cluster-contrast-reid)
================================================
FILE: examples/test.py
================================================
from __future__ import print_function, absolute_import
import argparse
import os.path as osp
import random
import numpy as np
import sys
import torch
from torch import nn
from torch.backends import cudnn
from torch.utils.data import DataLoader
from hhcl import datasets
from hhcl import models
from hhcl.models.dsbn import convert_dsbn, convert_bn
from hhcl.evaluators import Evaluator
from hhcl.utils.data import transforms as T
from hhcl.utils.data.preprocessor import Preprocessor
from hhcl.utils.logging import Logger
from hhcl.utils.serialization import load_checkpoint, save_checkpoint, copy_state_dict
def get_data(name, data_dir, height, width, batch_size, workers):
root = osp.join(data_dir, name)
dataset = datasets.create(name, root)
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.ToTensor(),
normalizer
])
test_loader = DataLoader(
Preprocessor(list(set(dataset.query) | set(dataset.gallery)),
root=dataset.images_dir, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
return dataset, test_loader
def main():
args = parser.parse_args()
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
main_worker(args)
def main_worker(args):
cudnn.benchmark = True
log_dir = osp.dirname(args.resume)
sys.stdout = Logger(osp.join(log_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create data loaders
dataset, test_loader = get_data(args.dataset, args.data_dir, args.height,
args.width, args.batch_size, args.workers)
# Create model
model = models.create(args.arch, pretrained=False, num_features=args.features, dropout=args.dropout,
num_classes=0, pooling_type=args.pooling_type)
if args.dsbn:
print("==> Load the model with domain-specific BNs")
convert_dsbn(model)
# Load from checkpoint
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model, strip='module.')
if args.dsbn:
print("==> Test with {}-domain BNs".format("source" if args.test_source else "target"))
convert_bn(model, use_target=(not args.test_source))
model.cuda()
model = nn.DataParallel(model)
# Evaluator
model.eval()
evaluator = Evaluator(model)
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=True, rerank=args.rerank)
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Testing the model")
# data
parser.add_argument('-d', '--dataset', type=str, default='market1501')
parser.add_argument('-b', '--batch-size', type=int, default=256)
parser.add_argument('-j', '--workers', type=int, default=4)
parser.add_argument('--height', type=int, default=256, help="input height")
parser.add_argument('--width', type=int, default=128, help="input width")
# model
parser.add_argument('-a', '--arch', type=str, default='resnet50',
choices=models.names())
parser.add_argument('--features', type=int, default=0)
parser.add_argument('--dropout', type=float, default=0)
parser.add_argument('--resume', type=str,
default="examples/logs/market1501/resnet50_avg/model_best.pth.tar",
metavar='PATH')
# testing configs
parser.add_argument('--rerank', action='store_true',
help="evaluation only")
parser.add_argument('--dsbn', action='store_true',
help="test on the model with domain-specific BN")
parser.add_argument('--test-source', action='store_true',
help="test on the source domain")
parser.add_argument('--seed', type=int, default=1)
# path
working_dir = osp.dirname(osp.abspath(__file__))
parser.add_argument('--data-dir', type=str, metavar='PATH',
default='examples/data')
parser.add_argument('--pooling-type', type=str, default='avg')
parser.add_argument('--embedding_features_path', type=str,
default='examples/logs/market1501/resnet50_avg/')
main()
================================================
FILE: examples/train.py
================================================
# -*- coding: utf-8 -*-
from __future__ import print_function, absolute_import
import argparse
import os.path as osp
import random
import numpy as np
import sys
import collections
import time
from datetime import timedelta
from sklearn.cluster import DBSCAN
import torch
from torch import nn
from torch.backends import cudnn
from torch.utils.data import DataLoader
import torch.nn.functional as F
from hhcl import datasets
from hhcl import models
from hhcl.models.cm import ClusterMemory
from hhcl.trainers import Trainer
from hhcl.evaluators import Evaluator, extract_features
from hhcl.utils.data import IterLoader
from hhcl.utils.data import transforms as T
from hhcl.utils.data.sampler import RandomMultipleGallerySampler
from hhcl.utils.data.preprocessor import Preprocessor
from hhcl.utils.logging import Logger
from hhcl.utils.serialization import load_checkpoint, save_checkpoint, copy_state_dict
from hhcl.utils.faiss_rerank import compute_jaccard_distance
start_epoch = best_mAP = 0
def get_data(name, data_dir):
root = osp.join(data_dir, name)
dataset = datasets.create(name, root)
return dataset
def get_train_loader(args, dataset, height, width, batch_size, workers,
num_instances, iters, trainset=None):
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.RandomHorizontalFlip(p=0.5),
T.Pad(10),
T.RandomCrop((height, width)),
T.ToTensor(),
normalizer,
T.RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406])
])
train_set = sorted(dataset.train) if trainset is None else sorted(trainset)
rmgs_flag = num_instances > 0
if rmgs_flag:
sampler = RandomMultipleGallerySampler(train_set, num_instances)
else:
sampler = None
train_loader = IterLoader(
DataLoader(Preprocessor(train_set, root=dataset.images_dir, transform=train_transformer),
batch_size=batch_size, num_workers=workers, sampler=sampler,
shuffle=not rmgs_flag, pin_memory=True, drop_last=True), length=iters)
return train_loader
def get_test_loader(dataset, height, width, batch_size, workers, testset=None):
normalizer = T.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_transformer = T.Compose([
T.Resize((height, width), interpolation=3),
T.ToTensor(),
normalizer
])
if testset is None:
testset = list(set(dataset.query) | set(dataset.gallery))
test_loader = DataLoader(
Preprocessor(testset, root=dataset.images_dir, transform=test_transformer),
batch_size=batch_size, num_workers=workers,
shuffle=False, pin_memory=True)
return test_loader
def create_model(args):
model = models.create(args.arch, num_features=args.features, norm=True, dropout=args.dropout,
num_classes=0, pooling_type=args.pooling_type)
# Load from checkpoint
if args.resume:
global start_epoch
checkpoint = load_checkpoint(args.resume)
copy_state_dict(checkpoint['state_dict'], model, strip='module.')
start_epoch = checkpoint['epoch']
# use CUDA
model.cuda()
model = nn.DataParallel(model)
return model
def main():
args = parser.parse_args()
if args.seed is not None:
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
main_worker(args)
def main_worker(args):
global start_epoch, best_mAP
start_time = time.monotonic()
cudnn.benchmark = True
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
print("==========\nArgs:{}\n==========".format(args))
# Create datasets
iters = args.iters if (args.iters > 0) else None
print("==> Load unlabeled dataset")
dataset = get_data(args.dataset, args.data_dir)
test_loader = get_test_loader(dataset, args.height, args.width, args.batch_size, args.workers)
# Create model
model = create_model(args)
# Evaluator
evaluator = Evaluator(model)
# Optimizer
params = [{"params": [value]} for _, value in model.named_parameters() if value.requires_grad]
optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=0.1)
# Trainer
trainer = Trainer(model)
for epoch in range(start_epoch, args.epochs):
with torch.no_grad():
print('==> Create pseudo labels for unlabeled data')
cluster_loader = get_test_loader(dataset, args.height, args.width,
args.batch_size, args.workers, testset=sorted(dataset.train))
features, _ = extract_features(model, cluster_loader, print_freq=50)
features = torch.cat([features[f].unsqueeze(0) for f, _, _ in sorted(dataset.train)], 0)
rerank_dist = compute_jaccard_distance(features, k1=args.k1, k2=args.k2)
if epoch == start_epoch:
# DBSCAN cluster
eps = args.eps
print('Clustering criterion eps: {:.3f}'.format(eps))
cluster = DBSCAN(eps=eps, min_samples=4, metric='precomputed', n_jobs=-1)
# select & cluster images as training set of this epochs
pseudo_labels = cluster.fit_predict(rerank_dist)
num_cluster = len(set(pseudo_labels)) - (1 if -1 in pseudo_labels else 0)
# generate new dataset and calculate cluster centers
@torch.no_grad()
def generate_cluster_features(labels, features):
centers = collections.defaultdict(list)
for i, label in enumerate(labels):
if label == -1:
continue
centers[labels[i]].append(features[i])
centers = [
torch.stack(centers[idx], dim=0).mean(0) for idx in sorted(centers.keys())
]
centers = torch.stack(centers, dim=0)
return centers
cluster_features = generate_cluster_features(pseudo_labels, features)
def generate_random_features(labels, features, num_cluster, num_instances):
indexes = np.zeros(num_cluster*num_instances)
for i in range(num_cluster):
index = [i+k*num_cluster for k in range(num_instances)]
samples = np.random.choice(np.where(pseudo_labels==i)[0], num_instances, True)
indexes[index] = samples
memory_features = features[indexes]
return memory_features
if args.memorybank=='CMhybrid_v2':
memory_features = generate_random_features(pseudo_labels, features, num_cluster, args.num_instances)
mask = (pseudo_labels < 0).astype(int)
print('==> Statistics for outliers with pseudo labels. outliers/total = {}/{} = {:.3f}'.format(mask.sum(), pseudo_labels.size, mask.sum()/pseudo_labels.size))
del cluster_loader, features
# Create memory bank
memory = ClusterMemory(model.module.num_features, num_cluster, temp=args.temp,
momentum=args.momentum, mode=args.memorybank, smooth=args.smooth,
num_instances=args.num_instances).cuda()
if args.memorybank=='CMhybrid':
memory.features = F.normalize(cluster_features.repeat(2, 1), dim=1).cuda()
elif args.memorybank=='CMhybrid_v2':
memory.features = F.normalize(torch.cat([cluster_features, memory_features],dim=0), dim=1).cuda()
else:
memory.features = F.normalize(cluster_features, dim=1).cuda()
trainer.memory = memory
pseudo_labeled_dataset = []
for i, ((fname, _, cid), label) in enumerate(zip(sorted(dataset.train), pseudo_labels)):
if label != -1:
pseudo_labeled_dataset.append((fname, label.item(), cid))
print('==> Statistics for epoch {}: {} clusters'.format(epoch, num_cluster))
train_loader = get_train_loader(args, dataset, args.height, args.width,
args.batch_size, args.workers, args.num_instances, iters,
trainset=pseudo_labeled_dataset)
train_loader.new_epoch()
trainer.train(epoch, train_loader, optimizer,
print_freq=args.print_freq, train_iters=len(train_loader))
if (epoch + 1) % args.eval_step == 0 or (epoch == args.epochs - 1):
mAP = evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=False)
is_best = (mAP > best_mAP)
best_mAP = max(mAP, best_mAP)
save_checkpoint({
'state_dict': model.state_dict(),
'epoch': epoch + 1,
'best_mAP': best_mAP,
}, is_best, fpath=osp.join(args.logs_dir, 'checkpoint.pth.tar'))
print('\n * Finished epoch {:3d} model mAP: {:5.1%} best: {:5.1%}{}\n'.
format(epoch, mAP, best_mAP, ' *' if is_best else ''))
lr_scheduler.step()
print('==> Test with the best model:')
checkpoint = load_checkpoint(osp.join(args.logs_dir, 'model_best.pth.tar'))
model.load_state_dict(checkpoint['state_dict'])
evaluator.evaluate(test_loader, dataset.query, dataset.gallery, cmc_flag=True)
end_time = time.monotonic()
print('Total running time: ', timedelta(seconds=end_time - start_time))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Hard-sample Guided Hybrid Contrast Learning for Unsupervised Person Re-ID")
# data
parser.add_argument('-d', '--dataset', type=str, default='dukemtmcreid',
choices=datasets.names())
parser.add_argument('-b', '--batch-size', type=int, default=2)
parser.add_argument('-j', '--workers', type=int, default=4)
parser.add_argument('--height', type=int, default=256, help="input height")
parser.add_argument('--width', type=int, default=128, help="input width")
parser.add_argument('--num-instances', type=int, default=4,
help="each minibatch consist of "
"(batch_size // num_instances) identities, and "
"each identity has num_instances instances, "
"default: 0 (NOT USE)")
# cluster
parser.add_argument('--eps', type=float, default=0.6,
help="max neighbor distance for DBSCAN")
parser.add_argument('--eps-gap', type=float, default=0.02,
help="multi-scale criterion for measuring cluster reliability")
parser.add_argument('--k1', type=int, default=30,
help="hyperparameter for jaccard distance")
parser.add_argument('--k2', type=int, default=6,
help="hyperparameter for jaccard distance")
# model
parser.add_argument('-a', '--arch', type=str, default='resnet50',
choices=models.names())
parser.add_argument('--features', type=int, default=0)
parser.add_argument('--dropout', type=float, default=0)
parser.add_argument('--smooth', type=float, default=0, help="label smoothing")
parser.add_argument('--hard-weight', type=float, default=0.5, help="hard weights")
parser.add_argument('--momentum', type=float, default=0.1,
help="update momentum for the memory bank")
parser.add_argument('--pooling-type', type=str, default='gem')
parser.add_argument('-mb', '--memorybank', type=str, default='CM', choices=['CM', 'CMhard', 'CMhybrid', 'CMhybrid_v2'])
# optimizer
parser.add_argument('--lr', type=float, default=0.00035,
help="learning rate")
parser.add_argument('--weight-decay', type=float, default=5e-4)
parser.add_argument('--epochs', type=int, default=50)
parser.add_argument('--iters', type=int, default=400)
parser.add_argument('--step-size', type=int, default=20)
# training configs
parser.add_argument('--seed', type=int, default=1)
parser.add_argument('--print-freq', type=int, default=10)
parser.add_argument('--eval-step', type=int, default=10)
parser.add_argument('--temp', type=float, default=0.05,
help="temperature for scaling contrastive loss")
# path
working_dir = osp.dirname(osp.abspath(__file__))
parser.add_argument('--data-dir', type=str, metavar='PATH',
default=osp.join(working_dir, 'data'))
parser.add_argument('--logs-dir', type=str, metavar='PATH',
default=osp.join(working_dir, 'logs'))
parser.add_argument('--resume', type=str, metavar='PATH', default='')
main()
================================================
FILE: hhcl/__init__.py
================================================
from __future__ import absolute_import
from . import datasets
from . import evaluation_metrics
from . import models
from . import utils
from . import evaluators
from . import trainers
__version__ = '0.1.0'
================================================
FILE: hhcl/datasets/__init__.py
================================================
from __future__ import absolute_import
import warnings
from .market1501 import Market1501
from .msmt17 import MSMT17
from .personx import PersonX
from .dukemtmcreid import DukeMTMCreID
from .celebreid import CelebReID
__factory = {
'market1501': Market1501,
'msmt17': MSMT17,
'personx': PersonX,
'dukemtmcreid': DukeMTMCreID,
'celebreid': CelebReID
}
def names():
return sorted(__factory.keys())
def create(name, root, *args, **kwargs):
"""
Create a dataset instance.
Parameters
----------
name : str
The dataset name.
root : str
The path to the dataset directory.
split_id : int, optional
The index of data split. Default: 0
num_val : int or float, optional
When int, it means the number of validation identities. When float,
it means the proportion of validation to all the trainval. Default: 100
download : bool, optional
If True, will download the dataset. Default: False
"""
if name not in __factory:
raise KeyError("Unknown dataset:", name)
return __factory[name](root, *args, **kwargs)
def get_dataset(name, root, *args, **kwargs):
warnings.warn("get_dataset is deprecated. Use create instead.")
return create(name, root, *args, **kwargs)
================================================
FILE: hhcl/datasets/celebreid.py
================================================
from __future__ import print_function, absolute_import
import os.path as osp
import glob
import re
from ..utils.data import BaseImageDataset
class CelebReID(BaseImageDataset):
"""
CelebReID
"""
dataset_dir = 'CelebReID'
def __init__(self, root, verbose=True, **kwargs):
super(CelebReID, self).__init__()
self.dataset_dir = osp.join(root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'gallery')
self._check_before_run()
train = self._process_dir(self.train_dir, relabel=True)
query = self._process_dir(self.query_dir, relabel=False)
gallery = self._process_dir(self.gallery_dir, relabel=False)
if verbose:
print("=> CelebReID loaded")
self.print_dataset_statistics(train, query, gallery)
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
def _check_before_run(self):
"""Check if all files are available before going deeper"""
if not osp.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
def _process_dir(self, dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
# pattern = re.compile(r'([-\d]+)_c(\d)')
pattern = re.compile(r'([-\d]+)_(\d)')
pid_container = set()
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
dataset = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
# assert 0 <= pid <= 1501 # pid == 0 means background
# assert 1 <= camid <= 6
camid -= 1 # index starts from 0
if relabel:
pid = pid2label[pid]
dataset.append((img_path, pid, camid))
return dataset
================================================
FILE: hhcl/datasets/dukemtmcreid.py
================================================
import glob
import os.path as osp
import re
from ..utils.data import BaseImageDataset
def process_dir(dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, "*.jpg"))
pattern = re.compile(r"([-\d]+)_c(\d)")
# get all identities
pid_container = set()
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
if pid == -1:
continue
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
data = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if (pid not in pid_container) or (pid == -1):
continue
assert 1 <= camid <= 8
camid -= 1
if relabel:
pid = pid2label[pid]
data.append((img_path, pid, camid))
return data
class DukeMTMCreID(BaseImageDataset):
"""DukeMTMC-reID.
Reference:
- Ristani et al. Performance Measures and a Data Set for Multi-Target,
Multi-Camera Tracking. ECCVW 2016.
- Zheng et al. Unlabeled Samples Generated by GAN Improve the Person
Re-identification Baseline in vitro. ICCV 2017.
URL: `<https://github.com/layumi/DukeMTMC-reID_evaluation>`_
Dataset statistics:
- identities: 1404 (train + query).
- images:16522 (train) + 2228 (query) + 17661 (gallery).
- cameras: 8.
"""
dataset_dir = "DukeMTMC-reID"
def __init__(self, root, verbose=True):
super(DukeMTMCreID, self).__init__()
self.root = osp.abspath(osp.expanduser(root))
self.dataset_dir = osp.join(self.root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
train = process_dir(dir_path=self.train_dir, relabel=True)
query = process_dir(dir_path=self.query_dir, relabel=False)
gallery = process_dir(dir_path=self.gallery_dir, relabel=False)
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
def _check_before_run(self):
"""Check if all files are available before going deeper"""
if not osp.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
================================================
FILE: hhcl/datasets/market1501.py
================================================
from __future__ import print_function, absolute_import
import os.path as osp
import glob
import re
from ..utils.data import BaseImageDataset
class Market1501(BaseImageDataset):
"""
Market1501
Reference:
Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
URL: http://www.liangzheng.org/Project/project_reid.html
Dataset statistics:
# identities: 1501 (+1 for background)
# images: 12936 (train) + 3368 (query) + 15913 (gallery)
"""
dataset_dir = 'Market-1501-v15.09.15'
def __init__(self, root, verbose=True, **kwargs):
super(Market1501, self).__init__()
self.dataset_dir = osp.join(root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
self._check_before_run()
train = self._process_dir(self.train_dir, relabel=True)
query = self._process_dir(self.query_dir, relabel=False)
gallery = self._process_dir(self.gallery_dir, relabel=False)
if verbose:
print("=> Market1501 loaded")
self.print_dataset_statistics(train, query, gallery)
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
def _check_before_run(self):
"""Check if all files are available before going deeper"""
if not osp.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
def _process_dir(self, dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
pattern = re.compile(r'([-\d]+)_c(\d)')
pid_container = set()
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
dataset = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
assert 0 <= pid <= 1501 # pid == 0 means background
assert 1 <= camid <= 6
camid -= 1 # index starts from 0
if relabel:
pid = pid2label[pid]
dataset.append((img_path, pid, camid))
return dataset
================================================
FILE: hhcl/datasets/msmt17.py
================================================
from __future__ import print_function, absolute_import
import os.path as osp
import glob
import re
from ..utils.data import BaseImageDataset
def _process_dir(dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
pattern = re.compile(r'([-\d]+)_c(\d+)')
pid_container = set()
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
dataset = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1:
continue # junk images are just ignored
assert 1 <= camid <= 15
camid -= 1 # index starts from 0
if relabel:
pid = pid2label[pid]
dataset.append((img_path, pid, camid))
return dataset
class MSMT17(BaseImageDataset):
dataset_dir = 'MSMT17_V1'
def __init__(self, root, verbose=True, **kwargs):
super(MSMT17, self).__init__()
self.dataset_dir = osp.join(root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
self._check_before_run()
train = _process_dir(self.train_dir, relabel=True)
query = _process_dir(self.query_dir, relabel=False)
gallery = _process_dir(self.gallery_dir, relabel=False)
if verbose:
print("=> MSMT17_V1 loaded")
self.print_dataset_statistics(train, query, gallery)
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
def _check_before_run(self):
"""Check if all files are available before going deeper"""
if not osp.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
================================================
FILE: hhcl/datasets/personx.py
================================================
from __future__ import print_function, absolute_import
import os.path as osp
import glob
import re
from ..utils.data import BaseImageDataset
class PersonX(BaseImageDataset):
"""
PersonX
Reference:
Sun et al. Dissecting Person Re-identification from the Viewpoint of Viewpoint. CVPR 2019.
Dataset statistics:
# identities: 1266
# images: 9840 (train) + 5136 (query) + 30816 (gallery)
"""
dataset_dir = 'PersonX'
def __init__(self, root, verbose=True, **kwargs):
super(PersonX, self).__init__()
self.dataset_dir = osp.join(root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
self._check_before_run()
train = self._process_dir(self.train_dir, relabel=True)
query = self._process_dir(self.query_dir, relabel=False)
gallery = self._process_dir(self.gallery_dir, relabel=False)
if verbose:
print("=> PersonX loaded")
self.print_dataset_statistics(train, query, gallery)
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
def _check_before_run(self):
"""Check if all files are available before going deeper"""
if not osp.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
def _process_dir(self, dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
pattern = re.compile(r'([-\d]+)_c([-\d]+)')
cam2label = {3: 1, 4: 2, 8: 3, 10: 4, 11: 5, 12: 6}
pid_container = set()
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
pid_container.add(pid)
pid2label = {pid: label for label, pid in enumerate(pid_container)}
dataset = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
assert (camid in cam2label.keys())
camid = cam2label[camid]
camid -= 1 # index starts from 0
if relabel: pid = pid2label[pid]
dataset.append((img_path, pid, camid))
return dataset
================================================
FILE: hhcl/evaluation_metrics/__init__.py
================================================
from __future__ import absolute_import
from .classification import accuracy
from .ranking import cmc, mean_ap
__all__ = [
'accuracy',
'cmc',
'mean_ap'
]
================================================
FILE: hhcl/evaluation_metrics/classification.py
================================================
from __future__ import absolute_import
import torch
from ..utils import to_torch
def accuracy(output, target, topk=(1,)):
with torch.no_grad():
output, target = to_torch(output), to_torch(target)
maxk = max(topk)
batch_size = target.size(0)
_, pred = output.topk(maxk, 1, True, True)
pred = pred.t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
ret = []
for k in topk:
correct_k = correct[:k].view(-1).float().sum(dim=0, keepdim=True)
ret.append(correct_k.mul_(1. / batch_size))
return ret
================================================
FILE: hhcl/evaluation_metrics/ranking.py
================================================
from __future__ import absolute_import
from collections import defaultdict
import numpy as np
from sklearn.metrics import average_precision_score
from ..utils import to_numpy
def _unique_sample(ids_dict, num):
mask = np.zeros(num, dtype=np.bool)
for _, indices in ids_dict.items():
i = np.random.choice(indices)
mask[i] = True
return mask
def cmc(distmat, query_ids=None, gallery_ids=None,
query_cams=None, gallery_cams=None, topk=100,
separate_camera_set=False,
single_gallery_shot=False,
first_match_break=False):
distmat = to_numpy(distmat)
m, n = distmat.shape
# Fill up default values
if query_ids is None:
query_ids = np.arange(m)
if gallery_ids is None:
gallery_ids = np.arange(n)
if query_cams is None:
query_cams = np.zeros(m).astype(np.int32)
if gallery_cams is None:
gallery_cams = np.ones(n).astype(np.int32)
# Ensure numpy array
query_ids = np.asarray(query_ids)
gallery_ids = np.asarray(gallery_ids)
query_cams = np.asarray(query_cams)
gallery_cams = np.asarray(gallery_cams)
# Sort and find correct matches
indices = np.argsort(distmat, axis=1)
matches = (gallery_ids[indices] == query_ids[:, np.newaxis])
# Compute CMC for each query
ret = np.zeros(topk)
num_valid_queries = 0
for i in range(m):
# Filter out the same id and same camera
valid = ((gallery_ids[indices[i]] != query_ids[i]) |
(gallery_cams[indices[i]] != query_cams[i]))
if separate_camera_set:
# Filter out samples from same camera
valid &= (gallery_cams[indices[i]] != query_cams[i])
if not np.any(matches[i, valid]): continue
if single_gallery_shot:
repeat = 10
gids = gallery_ids[indices[i][valid]]
inds = np.where(valid)[0]
ids_dict = defaultdict(list)
for j, x in zip(inds, gids):
ids_dict[x].append(j)
else:
repeat = 1
for _ in range(repeat):
if single_gallery_shot:
# Randomly choose one instance for each id
sampled = (valid & _unique_sample(ids_dict, len(valid)))
index = np.nonzero(matches[i, sampled])[0]
else:
index = np.nonzero(matches[i, valid])[0]
delta = 1. / (len(index) * repeat)
for j, k in enumerate(index):
if k - j >= topk: break
if first_match_break:
ret[k - j] += 1
break
ret[k - j] += delta
num_valid_queries += 1
if num_valid_queries == 0:
raise RuntimeError("No valid query")
return ret.cumsum() / num_valid_queries
def mean_ap(distmat, query_ids=None, gallery_ids=None,
query_cams=None, gallery_cams=None):
distmat = to_numpy(distmat)
m, n = distmat.shape
# Fill up default values
if query_ids is None:
query_ids = np.arange(m)
if gallery_ids is None:
gallery_ids = np.arange(n)
if query_cams is None:
query_cams = np.zeros(m).astype(np.int32)
if gallery_cams is None:
gallery_cams = np.ones(n).astype(np.int32)
# Ensure numpy array
query_ids = np.asarray(query_ids)
gallery_ids = np.asarray(gallery_ids)
query_cams = np.asarray(query_cams)
gallery_cams = np.asarray(gallery_cams)
# Sort and find correct matches
indices = np.argsort(distmat, axis=1)
matches = (gallery_ids[indices] == query_ids[:, np.newaxis])
# Compute AP for each query
aps = []
for i in range(m):
# Filter out the same id and same camera
valid = ((gallery_ids[indices[i]] != query_ids[i]) |
(gallery_cams[indices[i]] != query_cams[i]))
y_true = matches[i, valid]
y_score = -distmat[i][indices[i]][valid]
if not np.any(y_true): continue
aps.append(average_precision_score(y_true, y_score))
if len(aps) == 0:
raise RuntimeError("No valid query")
return np.mean(aps)
================================================
FILE: hhcl/evaluators.py
================================================
from __future__ import print_function, absolute_import
import time
import collections
from collections import OrderedDict
import numpy as np
import torch
import random
import copy
from .evaluation_metrics import cmc, mean_ap
from .utils.meters import AverageMeter
from .utils.rerank import re_ranking
from .utils import to_torch
def extract_cnn_feature(model, inputs):
inputs = to_torch(inputs).cuda()
outputs = model(inputs)
outputs = outputs.data.cpu()
return outputs
def extract_features(model, data_loader, print_freq=50):
model.eval()
batch_time = AverageMeter()
data_time = AverageMeter()
features = OrderedDict()
labels = OrderedDict()
end = time.time()
with torch.no_grad():
for i, (imgs, fnames, pids, _, _) in enumerate(data_loader):
data_time.update(time.time() - end)
outputs = extract_cnn_feature(model, imgs)
for fname, output, pid in zip(fnames, outputs, pids):
features[fname] = output
labels[fname] = pid
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % print_freq == 0:
print('Extract Features: [{}/{}]\t'
'Time {:.3f} ({:.3f})\t'
'Data {:.3f} ({:.3f})\t'
.format(i + 1, len(data_loader),
batch_time.val, batch_time.avg,
data_time.val, data_time.avg))
return features, labels
def pairwise_distance(features, query=None, gallery=None):
if query is None and gallery is None:
n = len(features)
x = torch.cat(list(features.values()))
x = x.view(n, -1)
dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True) * 2
dist_m = dist_m.expand(n, n) - 2 * torch.mm(x, x.t())
return dist_m
x = torch.cat([features[f].unsqueeze(0) for f, _, _ in query], 0)
y = torch.cat([features[f].unsqueeze(0) for f, _, _ in gallery], 0)
m, n = x.size(0), y.size(0)
x = x.view(m, -1)
y = y.view(n, -1)
dist_m = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + \
torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t()
dist_m.addmm_(1, -2, x, y.t())
return dist_m, x.numpy(), y.numpy()
def evaluate_all(query_features, gallery_features, distmat, query=None, gallery=None,
query_ids=None, gallery_ids=None,
query_cams=None, gallery_cams=None,
cmc_topk=(1, 5, 10), cmc_flag=False):
if query is not None and gallery is not None:
query_ids = [pid for _, pid, _ in query]
gallery_ids = [pid for _, pid, _ in gallery]
query_cams = [cam for _, _, cam in query]
gallery_cams = [cam for _, _, cam in gallery]
else:
assert (query_ids is not None and gallery_ids is not None
and query_cams is not None and gallery_cams is not None)
# Compute mean AP
mAP = mean_ap(distmat, query_ids, gallery_ids, query_cams, gallery_cams)
print('Mean AP: {:4.1%}'.format(mAP))
cmc_configs = {
'market1501': dict(separate_camera_set=False,
single_gallery_shot=False,
first_match_break=True),}
cmc_scores = {name: cmc(distmat, query_ids, gallery_ids,
query_cams, gallery_cams, **params)
for name, params in cmc_configs.items()}
print('CMC Scores:')
for k in cmc_topk:
print(' top-{:<4}{:12.1%}'.format(k, cmc_scores['market1501'][k-1]))
if (not cmc_flag):
return mAP
return cmc_scores['market1501'], mAP
class Evaluator(object):
def __init__(self, model):
super(Evaluator, self).__init__()
self.model = model
def evaluate(self, data_loader, query, gallery, cmc_flag=False, rerank=False):
features, _ = extract_features(self.model, data_loader)
distmat, query_features, gallery_features = pairwise_distance(features, query, gallery)
results = evaluate_all(query_features, gallery_features, distmat, query=query, gallery=gallery, cmc_flag=cmc_flag)
if (not rerank):
return results
print('Applying person re-ranking ...')
distmat_qq, _, _ = pairwise_distance(features, query, query)
distmat_gg, _, _ = pairwise_distance(features, gallery, gallery)
distmat = re_ranking(distmat.numpy(), distmat_qq.numpy(), distmat_gg.numpy())
return evaluate_all(query_features, gallery_features, distmat, query=query, gallery=gallery, cmc_flag=cmc_flag)
================================================
FILE: hhcl/models/__init__.py
================================================
from __future__ import absolute_import
from .resnet import *
from .resnet_ibn import *
__factory = {
'resnet18': resnet18,
'resnet34': resnet34,
'resnet50': resnet50,
'resnet101': resnet101,
'resnet152': resnet152,
'resnet_ibn50a': resnet_ibn50a,
'resnet_ibn101a': resnet_ibn101a,
}
def names():
return sorted(__factory.keys())
def create(name, *args, **kwargs):
"""
Create a model instance.
Parameters
----------
name : str
Model name. Can be one of 'inception', 'resnet18', 'resnet34',
'resnet50', 'resnet101', and 'resnet152'.
pretrained : bool, optional
Only applied for 'resnet*' models. If True, will use ImageNet pretrained
model. Default: True
cut_at_pooling : bool, optional
If True, will cut the model before the last global pooling layer and
ignore the remaining kwargs. Default: False
num_features : int, optional
If positive, will append a Linear layer after the global pooling layer,
with this number of output units, followed by a BatchNorm layer.
Otherwise these layers will not be appended. Default: 256 for
'inception', 0 for 'resnet*'
norm : bool, optional
If True, will normalize the feature to be unit L2-norm for each sample.
Otherwise will append a ReLU layer after the above Linear layer if
num_features > 0. Default: False
dropout : float, optional
If positive, will append a Dropout layer with this dropout rate.
Default: 0
num_classes : int, optional
If positive, will append a Linear layer at the end as the classifier
with this number of output units. Default: 0
"""
if name not in __factory:
raise KeyError("Unknown model:", name)
return __factory[name](*args, **kwargs)
================================================
FILE: hhcl/models/cm.py
================================================
import collections
import numpy as np
from abc import ABC
import torch
import torch.nn.functional as F
from torch import nn, autograd
from .losses import CrossEntropyLabelSmooth, FocalTopLoss
class CM(autograd.Function):
@staticmethod
def forward(ctx, inputs, targets, features, momentum):
ctx.features = features
ctx.momentum = momentum
ctx.save_for_backward(inputs, targets)
outputs = inputs.mm(ctx.features.t())
return outputs
@staticmethod
def backward(ctx, grad_outputs):
inputs, targets = ctx.saved_tensors
grad_inputs = None
if ctx.needs_input_grad[0]:
grad_inputs = grad_outputs.mm(ctx.features)
# momentum update
for x, y in zip(inputs, targets):
ctx.features[y] = ctx.momentum * ctx.features[y] + (1. - ctx.momentum) * x
ctx.features[y] /= ctx.features[y].norm()
return grad_inputs, None, None, None
def cm(inputs, indexes, features, momentum=0.5):
return CM.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device))
class CM_Hard(autograd.Function):
@staticmethod
def forward(ctx, inputs, targets, features, momentum):
ctx.features = features
ctx.momentum = momentum
ctx.save_for_backward(inputs, targets)
outputs = inputs.mm(ctx.features.t())
return outputs
@staticmethod
def backward(ctx, grad_outputs):
inputs, targets = ctx.saved_tensors
grad_inputs = None
if ctx.needs_input_grad[0]:
grad_inputs = grad_outputs.mm(ctx.features)
batch_centers = collections.defaultdict(list)
for instance_feature, index in zip(inputs, targets.tolist()):
batch_centers[index].append(instance_feature)
for index, features in batch_centers.items():
distances = []
for feature in features:
distance = feature.unsqueeze(0).mm(ctx.features[index].unsqueeze(0).t())[0][0]
distances.append(distance.cpu().numpy())
median = np.argmin(np.array(distances))
ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * features[median]
ctx.features[index] /= ctx.features[index].norm()
return grad_inputs, None, None, None
def cm_hard(inputs, indexes, features, momentum=0.5):
return CM_Hard.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device))
class CM_Hybrid(autograd.Function):
@staticmethod
def forward(ctx, inputs, targets, features, momentum):
ctx.features = features
ctx.momentum = momentum
ctx.save_for_backward(inputs, targets)
outputs = inputs.mm(ctx.features.t())
return outputs
@staticmethod
def backward(ctx, grad_outputs):
inputs, targets = ctx.saved_tensors
nums = len(ctx.features)//2
grad_inputs = None
if ctx.needs_input_grad[0]:
grad_inputs = grad_outputs.mm(ctx.features)
batch_centers = collections.defaultdict(list)
for instance_feature, index in zip(inputs, targets.tolist()):
batch_centers[index].append(instance_feature)
for index, features in batch_centers.items():
distances = []
for feature in features:
distance = feature.unsqueeze(0).mm(ctx.features[index].unsqueeze(0).t())[0][0]
distances.append(distance.cpu().numpy())
median = np.argmin(np.array(distances))
ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * features[median]
ctx.features[index] /= ctx.features[index].norm()
mean = torch.stack(features, dim=0).mean(0)
ctx.features[index+nums] = ctx.features[index+nums] * ctx.momentum + (1 - ctx.momentum) * mean
ctx.features[index+nums] /= ctx.features[index+nums].norm()
return grad_inputs, None, None, None
def cm_hybrid(inputs, indexes, features, momentum=0.5):
return CM_Hybrid.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device))
class CM_Hybrid_v2(autograd.Function):
@staticmethod
def forward(ctx, inputs, targets, features, momentum, num_instances):
ctx.features = features
ctx.momentum = momentum
ctx.num_instances = num_instances
ctx.save_for_backward(inputs, targets)
outputs = inputs.mm(ctx.features.t())
return outputs
@staticmethod
def backward(ctx, grad_outputs):
inputs, targets = ctx.saved_tensors
nums = len(ctx.features)//(ctx.num_instances + 1)
grad_inputs = None
if ctx.needs_input_grad[0]:
grad_inputs = grad_outputs.mm(ctx.features)
batch_centers = collections.defaultdict(list)
updated = set()
for k, (instance_feature, index) in enumerate(zip(inputs, targets.tolist())):
batch_centers[index].append(instance_feature)
if index not in updated:
indexes = [index + nums*i for i in range(1, (targets==index).sum()+1)]
ctx.features[indexes] = inputs[targets==index]
# ctx.features[indexes] = ctx.features[indexes] * ctx.momentum + (1 - ctx.momentum) * inputs[targets==index]
# ctx.features[indexes] /= ctx.features[indexes].norm(dim=1, keepdim=True)
updated.add(index)
for index, features in batch_centers.items():
mean = torch.stack(features, dim=0).mean(0)
ctx.features[index] = ctx.features[index] * ctx.momentum + (1 - ctx.momentum) * mean
ctx.features[index] /= ctx.features[index].norm()
return grad_inputs, None, None, None, None
def cm_hybrid_v2(inputs, indexes, features, momentum=0.5, num_instances=16, *args):
return CM_Hybrid_v2.apply(inputs, indexes, features, torch.Tensor([momentum]).to(inputs.device), num_instances)
class ClusterMemory(nn.Module, ABC):
__CMfactory = {
'CM': cm,
'CMhard':cm_hard,
}
def __init__(self, num_features, num_samples, temp=0.05, momentum=0.2, mode='CM', hard_weight=0.5, smooth=0., num_instances=1):
super(ClusterMemory, self).__init__()
self.num_features = num_features
self.num_samples = num_samples
self.momentum = momentum
self.temp = temp
self.cm_type = mode
if smooth>0:
self.cross_entropy = CrossEntropyLabelSmooth(self.num_samples, 0.1, True)
print('>>> Using CrossEntropy with Label Smoothing.')
else:
self.cross_entropy = nn.CrossEntropyLoss().cuda()
if self.cm_type in ['CM', 'CMhard']:
self.register_buffer('features', torch.zeros(num_samples, num_features))
elif self.cm_type=='CMhybrid':
self.hard_weight = hard_weight
print('hard_weight: {}'.format(self.hard_weight))
self.register_buffer('features', torch.zeros(2*num_samples, num_features))
elif self.cm_type=='CMhybrid_v2':
self.hard_weight = hard_weight
self.num_instances = num_instances
self.register_buffer('features', torch.zeros((self.num_instances+1)*num_samples, num_features))
else:
raise TypeError('Cluster Memory {} is invalid!'.format(self.cm_type))
def forward(self, inputs, targets):
if self.cm_type in ['CM', 'CMhard']:
outputs = ClusterMemory.__CMfactory[self.cm_type](inputs, targets, self.features, self.momentum)
outputs /= self.temp
loss = self.cross_entropy(outputs, targets)
return loss
elif self.cm_type=='CMhybrid':
outputs = cm_hybrid(inputs, targets, self.features, self.momentum)
outputs /= self.temp
output_hard, output_mean = torch.chunk(outputs, 2, dim=1)
loss = self.hard_weight * (self.cross_entropy(output_hard, targets) + (1 - self.hard_weight) * self.cross_entropy(output_mean, targets))
return loss
elif self.cm_type=='CMhybrid_v2':
outputs = cm_hybrid_v2(inputs, targets, self.features, self.momentum, self.num_instances)
out_list = torch.chunk(outputs, self.num_instances+1, dim=1)
out = torch.stack(out_list[1:], dim=0)
neg = torch.max(out, dim=0)[0]
pos = torch.min(out, dim=0)[0]
mask = torch.zeros_like(out_list[0]).scatter_(1, targets.unsqueeze(1), 1)
logits = mask * pos + (1-mask) * neg
loss = self.hard_weight * self.cross_entropy(out_list[0]/self.temp, targets) \
+ (1 - self.hard_weight) * self.cross_entropy(logits/self.temp, targets)
return loss
================================================
FILE: hhcl/models/dsbn.py
================================================
import torch
import torch.nn as nn
# Domain-specific BatchNorm
class DSBN2d(nn.Module):
def __init__(self, planes):
super(DSBN2d, self).__init__()
self.num_features = planes
self.BN_S = nn.BatchNorm2d(planes)
self.BN_T = nn.BatchNorm2d(planes)
def forward(self, x):
if (not self.training):
return self.BN_T(x)
bs = x.size(0)
assert (bs%2==0)
split = torch.split(x, int(bs/2), 0)
out1 = self.BN_S(split[0].contiguous())
out2 = self.BN_T(split[1].contiguous())
out = torch.cat((out1, out2), 0)
return out
class DSBN1d(nn.Module):
def __init__(self, planes):
super(DSBN1d, self).__init__()
self.num_features = planes
self.BN_S = nn.BatchNorm1d(planes)
self.BN_T = nn.BatchNorm1d(planes)
def forward(self, x):
if (not self.training):
return self.BN_T(x)
bs = x.size(0)
assert (bs%2==0)
split = torch.split(x, int(bs/2), 0)
out1 = self.BN_S(split[0].contiguous())
out2 = self.BN_T(split[1].contiguous())
out = torch.cat((out1, out2), 0)
return out
def convert_dsbn(model):
for _, (child_name, child) in enumerate(model.named_children()):
assert(not next(model.parameters()).is_cuda)
if isinstance(child, nn.BatchNorm2d):
m = DSBN2d(child.num_features)
m.BN_S.load_state_dict(child.state_dict())
m.BN_T.load_state_dict(child.state_dict())
setattr(model, child_name, m)
elif isinstance(child, nn.BatchNorm1d):
m = DSBN1d(child.num_features)
m.BN_S.load_state_dict(child.state_dict())
m.BN_T.load_state_dict(child.state_dict())
setattr(model, child_name, m)
else:
convert_dsbn(child)
def convert_bn(model, use_target=True):
for _, (child_name, child) in enumerate(model.named_children()):
assert(not next(model.parameters()).is_cuda)
if isinstance(child, DSBN2d):
m = nn.BatchNorm2d(child.num_features)
if use_target:
m.load_state_dict(child.BN_T.state_dict())
else:
m.load_state_dict(child.BN_S.state_dict())
setattr(model, child_name, m)
elif isinstance(child, DSBN1d):
m = nn.BatchNorm1d(child.num_features)
if use_target:
m.load_state_dict(child.BN_T.state_dict())
else:
m.load_state_dict(child.BN_S.state_dict())
setattr(model, child_name, m)
else:
convert_bn(child, use_target=use_target)
================================================
FILE: hhcl/models/kmeans.py
================================================
# Written by Yixiao Ge
import warnings
import faiss
import torch
from ..utils import to_numpy, to_torch
__all__ = ["label_generator_kmeans"]
@torch.no_grad()
def label_generator_kmeans(features, num_classes=500, cuda=True):
assert num_classes, "num_classes for kmeans is null"
# k-means cluster by faiss
cluster = faiss.Kmeans(
features.size(-1), num_classes, niter=300, verbose=True, gpu=cuda
)
cluster.train(to_numpy(features))
_, labels = cluster.index.search(to_numpy(features), 1)
labels = labels.reshape(-1)
centers = to_torch(cluster.centroids).float()
# labels = to_torch(labels).long()
# k-means does not have outlier points
assert not (-1 in labels)
return labels, centers, num_classes, None
================================================
FILE: hhcl/models/losses.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter
class CrossEntropyLabelSmooth(nn.Module):
"""Cross entropy loss with label smoothing regularizer.
Reference:
Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
Equation: y = (1 - epsilon) * y + epsilon / K.
Args:
num_classes (int): number of classes.
epsilon (float): weight.
"""
def __init__(self, num_classes=0, epsilon=0.1, topk_smoothing=False):
super(CrossEntropyLabelSmooth, self).__init__()
self.num_classes = num_classes
self.epsilon = epsilon
self.logsoftmax = nn.LogSoftmax(dim=1).cuda()
self.k = 1 if not topk_smoothing else self.num_classes//50
def forward(self, inputs, targets):
"""
Args:
inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
targets: ground truth labels with shape (num_classes)
"""
log_probs = self.logsoftmax(inputs)
if self.k >1:
topk = torch.argsort(-log_probs)[:,:self.k]
targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1 - self.epsilon)
targets += torch.zeros_like(log_probs).scatter_(1, topk, self.epsilon / self.k)
else:
targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1)
targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
loss = (- targets * log_probs).mean(0).sum()
return loss
class SoftEntropy(nn.Module):
def __init__(self, input_prob=False):
super(SoftEntropy, self).__init__()
self.input_prob = input_prob
self.logsoftmax = nn.LogSoftmax(dim=1).cuda()
def forward(self, inputs, targets):
log_probs = self.logsoftmax(inputs)
if self.input_prob:
loss = (- targets.detach() * log_probs).mean(0).sum()
else:
loss = (- F.softmax(targets, dim=1).detach() * log_probs).mean(0).sum()
return loss
class SoftEntropySmooth(nn.Module):
def __init__(self, epsilon=0.1):
super(SoftEntropySmooth, self).__init__()
self.epsilon = epsilon
self.logsoftmax = nn.LogSoftmax(dim=1).cuda()
def forward(self, inputs, soft_targets, targets):
log_probs = self.logsoftmax(inputs)
targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1)
soft_targets = F.softmax(soft_targets, dim=1)
smooth_targets = (1 - self.epsilon) * targets + self.epsilon * soft_targets
loss = (- smooth_targets.detach() * log_probs).mean(0).sum()
return loss
class Softmax(nn.Module):
def __init__(self, feat_dim, num_class, temp=0.05):
super(Softmax, self).__init__()
self.weight = Parameter(torch.Tensor(feat_dim, num_class))
self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
self.temp = temp
def forward(self, feats, labels):
kernel_norm = F.normalize(self.weight, dim=0)
feats = F.normalize(feats)
outputs = feats.mm(kernel_norm)
outputs /= self.temp
loss = F.cross_entropy(outputs, labels)
return loss
class CircleLoss(nn.Module):
"""Implementation for "Circle Loss: A Unified Perspective of Pair Similarity Optimization"
Note: this is the classification based implementation of circle loss.
"""
def __init__(self, feat_dim, num_class, margin=0.25, gamma=256):
super(CircleLoss, self).__init__()
self.weight = Parameter(torch.Tensor(feat_dim, num_class))
self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
self.margin = margin
self.gamma = gamma
self.O_p = 1 + margin
self.O_n = -margin
self.delta_p = 1-margin
self.delta_n = margin
def forward(self, feats, labels):
kernel_norm = F.normalize(self.weight, dim=0)
feats = F.normalize(feats)
cos_theta = torch.mm(feats, kernel_norm)
cos_theta = cos_theta.clamp(-1, 1)
index_pos = torch.zeros_like(cos_theta)
index_pos.scatter_(1, labels.data.view(-1, 1), 1)
index_pos = index_pos.bool()
index_neg = torch.ones_like(cos_theta)
index_neg.scatter_(1, labels.data.view(-1, 1), 0)
index_neg = index_neg.bool()
alpha_p = torch.clamp_min(self.O_p - cos_theta.detach(), min=0.)
alpha_n = torch.clamp_min(cos_theta.detach() - self.O_n, min=0.)
logit_p = alpha_p * (cos_theta - self.delta_p)
logit_n = alpha_n * (cos_theta - self.delta_n)
output = cos_theta * 1.0
output[index_pos] = logit_p[index_pos]
output[index_neg] = logit_n[index_neg]
output *= self.gamma
return F.cross_entropy(output, labels)
class CosFace(nn.Module):
r"""Implement of CosFace (https://arxiv.org/pdf/1801.09414.pdf):
Args:
in_features: size of each input sample
out_features: size of each output sample
s: norm of input feature
m: margin
cos(theta)-m
"""
def __init__(self, feat_dim, num_class, s = 64.0, m = 0.35):
super(CosFace, self).__init__()
self.in_features = feat_dim
self.out_features = num_class
self.s = s
self.m = m
self.weight = Parameter(torch.FloatTensor(feat_dim, num_class))
nn.init.xavier_uniform_(self.weight)
def forward(self, input, label):
# --------------------------- cos(theta) & phi(theta) ---------------------------
# cosine = F.linear(F.normalize(input), F.normalize(self.weight, dim=1))
cosine = torch.mm(F.normalize(input), F.normalize(self.weight, dim=0))
phi = cosine - self.m
# --------------------------- convert label to one-hot ---------------------------
one_hot = torch.zeros(cosine.size(), device = 'cuda')
# one_hot = one_hot.cuda() if cosine.is_cuda else one_hot
one_hot.scatter_(1, label.view(-1, 1).long(), 1)
# -------------torch.where(out_i = {x_i if condition_i else y_i) -------------
output = (one_hot * phi) + ((1.0 - one_hot) * cosine) # you can use torch.where if your torch.__version__ is 0.4
output *= self.s
return F.cross_entropy(output, label)
def __repr__(self):
return self.__class__.__name__ + '(' \
+ 'in_features = ' + str(self.in_features) \
+ ', out_features = ' + str(self.out_features) \
+ ', s = ' + str(self.s) \
+ ', m = ' + str(self.m) + ')'
import math
class InstanceLoss(nn.Module):
def __init__(self, batch_size, temperature, device):
super(InstanceLoss, self).__init__()
self.batch_size = batch_size
self.temperature = temperature
self.device = device
self.mask = self.mask_correlated_samples(batch_size)
self.criterion = nn.CrossEntropyLoss(reduction="sum")
def mask_correlated_samples(self, batch_size):
N = 2 * batch_size
mask = torch.ones((N, N))
mask = mask.fill_diagonal_(0)
for i in range(batch_size):
mask[i, batch_size + i] = 0
mask[batch_size + i, i] = 0
mask = mask.bool()
return mask
def forward(self, z_i, z_j):
N = 2 * self.batch_size
z = torch.cat((z_i, z_j), dim=0)
sim = torch.matmul(z, z.T) / self.temperature
sim_i_j = torch.diag(sim, self.batch_size)
sim_j_i = torch.diag(sim, -self.batch_size)
positive_samples = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(N, 1)
negative_samples = sim[self.mask].reshape(N, -1)
labels = torch.zeros(N).to(positive_samples.device).long()
logits = torch.cat((positive_samples, negative_samples), dim=1)
loss = self.criterion(logits, labels)
loss /= N
return loss
class ClusterLoss(nn.Module):
def __init__(self, class_num, temperature, device):
super(ClusterLoss, self).__init__()
self.class_num = class_num
self.temperature = temperature
self.device = device
self.mask = self.mask_correlated_clusters(class_num)
self.criterion = nn.CrossEntropyLoss(reduction="sum")
self.similarity_f = nn.CosineSimilarity(dim=2)
def mask_correlated_clusters(self, class_num):
N = 2 * class_num
mask = torch.ones((N, N))
mask = mask.fill_diagonal_(0)
for i in range(class_num):
mask[i, class_num + i] = 0
mask[class_num + i, i] = 0
mask = mask.bool()
return mask
def forward(self, c_i, c_j):
p_i = c_i.sum(0).view(-1)
p_i /= p_i.sum()
ne_i = math.log(p_i.size(0)) + (p_i * torch.log(p_i)).sum()
p_j = c_j.sum(0).view(-1)
p_j /= p_j.sum()
ne_j = math.log(p_j.size(0)) + (p_j * torch.log(p_j)).sum()
ne_loss = ne_i + ne_j
c_i = c_i.t()
c_j = c_j.t()
N = 2 * self.class_num
c = torch.cat((c_i, c_j), dim=0)
sim = self.similarity_f(c.unsqueeze(1), c.unsqueeze(0)) / self.temperature
sim_i_j = torch.diag(sim, self.class_num)
sim_j_i = torch.diag(sim, -self.class_num)
positive_clusters = torch.cat((sim_i_j, sim_j_i), dim=0).reshape(N, 1)
negative_clusters = sim[self.mask].reshape(N, -1)
labels = torch.zeros(N).to(positive_clusters.device).long()
logits = torch.cat((positive_clusters, negative_clusters), dim=1)
loss = self.criterion(logits, labels)
loss /= N
return loss + ne_loss
class FocalLoss(nn.Module):
def __init__(self, gamma=2, alpha=0.25):
super(FocalLoss, self).__init__()
self.alpha = alpha
self.gamma = gamma
print('Initializing FocalLoss for training: alpha={}, gamma={}'.format(self.alpha, self.gamma))
def forward(self, input, target):
assert input.dim() == 2
assert not target.requires_grad
target = target.squeeze(1) if target.dim() == 2 else target
assert target.dim() == 1
logpt = F.log_softmax(input, dim=1)
logpt_gt = logpt.gather(1,target.unsqueeze(1))
logpt_gt = logpt_gt.view(-1)
pt_gt = logpt_gt.exp()
assert logpt_gt.size() == pt_gt.size()
loss = -self.alpha*(torch.pow((1-pt_gt), self.gamma))*logpt_gt
return loss.mean()
class LabelRefineLoss(nn.Module):
def __init__(self, lambda1=0.0):
super(LabelRefineLoss, self).__init__()
self.lambda1 = lambda1
print('Initializing LabelRefineLoss for training: lambda1={}'.format(self.lambda1))
def forward(self, input, target):
assert input.dim() == 2
assert not target.requires_grad
target = target.squeeze(1) if target.dim() == 2 else target
assert target.dim() == 1
logpt = F.log_softmax(input, dim=1)
logpt_gt = logpt.gather(1,target.unsqueeze(1))
logpt_gt = logpt_gt.view(-1)
logpt_pred,_ = torch.max(logpt,1)
logpt_pred = logpt_pred.view(-1)
assert logpt_gt.size() == logpt_pred.size()
loss = - (1-self.lambda1)*logpt_gt - self.lambda1* logpt_pred
return loss.mean()
class FocalTopLoss(nn.Module):
def __init__(self, top_percent=0.7):
super(FocalTopLoss, self).__init__()
self.top_percent = top_percent
def masked_softmax_multi_focal(self, vec, targets=None, dim=1):
exps = torch.exp(vec)
one_hot_pos = F.one_hot(targets, num_classes=exps.shape[1])
one_hot_neg = one_hot_pos.new_ones(size=one_hot_pos.shape)
one_hot_neg = one_hot_neg - one_hot_pos
neg_exps = exps.new_zeros(size=exps.shape)
neg_exps[one_hot_neg>0] = exps[one_hot_neg>0]
ori_neg_exps = neg_exps
neg_exps = neg_exps/neg_exps.sum(dim=1, keepdim=True)
new_exps = exps.new_zeros(size=exps.shape)
new_exps[one_hot_pos>0] = exps[one_hot_pos>0]
sorted, indices = torch.sort(neg_exps, dim=1, descending=True)
sorted_cum_sum = torch.cumsum(sorted, dim=1)
sorted_cum_diff = (sorted_cum_sum - self.top_percent).abs()
sorted_cum_min_indices = sorted_cum_diff.argmin(dim=1)
min_values = sorted[torch.range(0, sorted.shape[0]-1).long(), sorted_cum_min_indices]
min_values = min_values.unsqueeze(dim=-1) * ori_neg_exps.sum(dim=1, keepdim=True)
ori_neg_exps[ori_neg_exps<min_values] = 0
new_exps[one_hot_neg>0] = ori_neg_exps[one_hot_neg>0]
masked_sums = exps.sum(dim, keepdim=True)
return new_exps / masked_sums
def forward(self, input, target):
masked_sim = self.masked_softmax_multi_focal(input, target)
return F.nll_loss(torch.log(masked_sim + 1e-6), target)
================================================
FILE: hhcl/models/pooling.py
================================================
# Credit to https://github.com/JDAI-CV/fast-reid/blob/master/fastreid/layers/pooling.py
from abc import ABC
import torch
import torch.nn.functional as F
from torch import nn
__all__ = [
"GeneralizedMeanPoolingPFpn",
"GeneralizedMeanPoolingList",
"GeneralizedMeanPoolingP",
"AdaptiveAvgMaxPool2d",
"FastGlobalAvgPool2d",
"avg_pooling",
"max_pooling",
]
class GeneralizedMeanPoolingList(nn.Module, ABC):
r"""Applies a 2D power-average adaptive pooling over an input signal composed of
several input planes.
The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)`
- At p = infinity, one gets Max Pooling
- At p = 1, one gets Average Pooling
The output is of size H x W, for any input size.
The number of output features is equal to the number of input planes.
Args:
output_size: the target output size of the image of the form H x W.
Can be a tuple (H, W) or a single H for a square image H x H
H and W can be either a ``int``, or ``None`` which means the size
will be the same as that of the input.
"""
def __init__(self, output_size=1, eps=1e-6):
super(GeneralizedMeanPoolingList, self).__init__()
self.output_size = output_size
self.eps = eps
def forward(self, x_list):
outs = []
for x in x_list:
x = x.clamp(min=self.eps)
out = torch.nn.functional.adaptive_avg_pool2d(x, self.output_size)
outs.append(out)
return torch.stack(outs, -1).mean(-1)
def __repr__(self):
return (
self.__class__.__name__
+ "("
+ "output_size="
+ str(self.output_size)
+ ")"
)
class GeneralizedMeanPooling(nn.Module, ABC):
r"""Applies a 2D power-average adaptive pooling over an input signal composed of
several input planes.
The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)`
- At p = infinity, one gets Max Pooling
- At p = 1, one gets Average Pooling
The output is of size H x W, for any input size.
The number of output features is equal to the number of input planes.
Args:
output_size: the target output size of the image of the form H x W.
Can be a tuple (H, W) or a single H for a square image H x H
H and W can be either a ``int``, or ``None`` which means the size
will be the same as that of the input.
"""
def __init__(self, norm, output_size=1, eps=1e-6):
super(GeneralizedMeanPooling, self).__init__()
assert norm > 0
self.p = float(norm)
self.output_size = output_size
self.eps = eps
def forward(self, x):
x = x.clamp(min=self.eps).pow(self.p)
return torch.nn.functional.adaptive_avg_pool2d(x, self.output_size).pow(
1.0 / self.p
)
def __repr__(self):
return (
self.__class__.__name__
+ "("
+ str(self.p)
+ ", "
+ "output_size="
+ str(self.output_size)
+ ")"
)
class GeneralizedMeanPoolingP(GeneralizedMeanPooling, ABC):
""" Same, but norm is trainable
"""
def __init__(self, norm=3, output_size=1, eps=1e-6):
super(GeneralizedMeanPoolingP, self).__init__(norm, output_size, eps)
self.p = nn.Parameter(torch.ones(1) * norm)
class GeneralizedMeanPoolingFpn(nn.Module, ABC):
r"""Applies a 2D power-average adaptive pooling over an input signal composed of
several input planes.
The function computed is: :math:`f(X) = pow(sum(pow(X, p)), 1/p)`
- At p = infinity, one gets Max Pooling
- At p = 1, one gets Average Pooling
The output is of size H x W, for any input size.
The number of output features is equal to the number of input planes.
Args:
output_size: the target output size of the image of the form H x W.
Can be a tuple (H, W) or a single H for a square image H x H
H and W can be either a ``int``, or ``None`` which means the size
will be the same as that of the input.
"""
def __init__(self, norm, output_size=1, eps=1e-6):
super(GeneralizedMeanPoolingFpn, self).__init__()
assert norm > 0
self.p = float(norm)
self.output_size = output_size
self.eps = eps
def forward(self, x_lists):
outs = []
for x in x_lists:
x = x.clamp(min=self.eps).pow(self.p)
out = torch.nn.functional.adaptive_avg_pool2d(x, self.output_size).pow(
1.0 / self.p
)
outs.append(out)
return torch.cat(outs, 1)
def __repr__(self):
return (
self.__class__.__name__
+ "("
+ str(self.p)
+ ", "
+ "output_size="
+ str(self.output_size)
+ ")"
)
class GeneralizedMeanPoolingPFpn(GeneralizedMeanPoolingFpn, ABC):
""" Same, but norm is trainable
"""
def __init__(self, norm=3, output_size=1, eps=1e-6):
super(GeneralizedMeanPoolingPFpn, self).__init__(norm, output_size, eps)
self.p = nn.Parameter(torch.ones(1) * norm)
class AdaptiveAvgMaxPool2d(nn.Module, ABC):
def __init__(self):
super(AdaptiveAvgMaxPool2d, self).__init__()
self.avgpool = FastGlobalAvgPool2d()
def forward(self, x):
x_avg = self.avgpool(x, self.output_size)
x_max = F.adaptive_max_pool2d(x, 1)
x = x_max + x_avg
return x
class FastGlobalAvgPool2d(nn.Module, ABC):
def __init__(self, flatten=False):
super(FastGlobalAvgPool2d, self).__init__()
self.flatten = flatten
def forward(self, x):
if self.flatten:
in_size = x.size()
return x.view((in_size[0], in_size[1], -1)).mean(dim=2)
else:
return (
x.view(x.size(0), x.size(1), -1)
.mean(-1)
.view(x.size(0), x.size(1), 1, 1)
)
def avg_pooling():
return nn.AdaptiveAvgPool2d(1)
# return FastGlobalAvgPool2d()
def max_pooling():
return nn.AdaptiveMaxPool2d(1)
class Flatten(nn.Module):
def forward(self, input):
return input.view(input.size(0), -1)
__pooling_factory = {
"avg": avg_pooling,
"max": max_pooling,
"gem": GeneralizedMeanPoolingP,
"gemFpn": GeneralizedMeanPoolingPFpn,
"gemList": GeneralizedMeanPoolingList,
"avg+max": AdaptiveAvgMaxPool2d,
}
def pooling_names():
return sorted(__pooling_factory.keys())
def build_pooling_layer(name):
"""
Create a pooling layer.
Parameters
----------
name : str
The backbone name.
"""
if name not in __pooling_factory:
raise KeyError("Unknown pooling layer:", name)
return __pooling_factory[name]()
================================================
FILE: hhcl/models/resnet.py
================================================
from __future__ import absolute_import
from torch import nn
from torch.nn import functional as F
from torch.nn import init
import torchvision
import torch
from .pooling import build_pooling_layer
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
'resnet152']
class ResNet(nn.Module):
__factory = {
18: torchvision.models.resnet18,
34: torchvision.models.resnet34,
50: torchvision.models.resnet50,
101: torchvision.models.resnet101,
152: torchvision.models.resnet152,
}
def __init__(self, depth, pretrained=True, cut_at_pooling=False,
num_features=0, norm=False, dropout=0, num_classes=0, pooling_type='avg'):
print('pooling_type: {}'.format(pooling_type))
super(ResNet, self).__init__()
self.pretrained = pretrained
self.depth = depth
self.cut_at_pooling = cut_at_pooling
# Construct base (pretrained) resnet
if depth not in ResNet.__factory:
raise KeyError("Unsupported depth:", depth)
resnet = ResNet.__factory[depth](pretrained=pretrained)
if self.depth >= 50:
resnet.layer4[0].conv2.stride = (1, 1)
resnet.layer4[0].downsample[0].stride = (1, 1)
self.base = nn.Sequential(
resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4)
self.gap = build_pooling_layer(pooling_type)
if not self.cut_at_pooling:
self.num_features = num_features
self.norm = norm
self.dropout = dropout
self.has_embedding = num_features > 0
self.num_classes = num_classes
out_planes = resnet.fc.in_features
# Append new layers
if self.has_embedding:
self.feat = nn.Linear(out_planes, self.num_features)
self.feat_bn = nn.BatchNorm1d(self.num_features)
init.kaiming_normal_(self.feat.weight, mode='fan_out')
init.constant_(self.feat.bias, 0)
else:
# Change the num_features to CNN output channels
self.num_features = out_planes
self.feat_bn = nn.BatchNorm1d(self.num_features)
self.feat_bn.bias.requires_grad_(False)
if self.dropout > 0:
self.drop = nn.Dropout(self.dropout)
if self.num_classes > 0:
self.classifier = nn.Linear(self.num_features, self.num_classes, bias=False)
init.normal_(self.classifier.weight, std=0.001)
init.constant_(self.feat_bn.weight, 1)
init.constant_(self.feat_bn.bias, 0)
if not pretrained:
self.reset_params()
def forward(self, x):
bs = x.size(0)
x = self.base(x)
x = self.gap(x)
x = x.view(x.size(0), -1)
if self.cut_at_pooling:
return x
if self.has_embedding:
bn_x = self.feat_bn(self.feat(x))
else:
bn_x = self.feat_bn(x)
if (self.training is False):
bn_x = F.normalize(bn_x)
return bn_x
if self.norm:
bn_x = F.normalize(bn_x)
elif self.has_embedding:
bn_x = F.relu(bn_x)
if self.dropout > 0:
bn_x = self.drop(bn_x)
if self.num_classes > 0:
prob = self.classifier(bn_x)
else:
return bn_x
return prob, bn_x
def reset_params(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm1d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
def resnet18(**kwargs):
return ResNet(18, **kwargs)
def resnet34(**kwargs):
return ResNet(34, **kwargs)
def resnet50(**kwargs):
return ResNet(50, **kwargs)
def resnet101(**kwargs):
return ResNet(101, **kwargs)
def resnet152(**kwargs):
return ResNet(152, **kwargs)
================================================
FILE: hhcl/models/resnet_ibn.py
================================================
from __future__ import absolute_import
from torch import nn
from torch.nn import functional as F
from torch.nn import init
import torchvision
import torch
from .pooling import build_pooling_layer
from .resnet_ibn_a import resnet50_ibn_a, resnet101_ibn_a
__all__ = ['ResNetIBN', 'resnet_ibn50a', 'resnet_ibn101a']
class ResNetIBN(nn.Module):
__factory = {
'50a': resnet50_ibn_a,
'101a': resnet101_ibn_a
}
def __init__(self, depth, pretrained=True, cut_at_pooling=False,
num_features=0, norm=False, dropout=0, num_classes=0, pooling_type='avg'):
print('pooling_type: {}'.format(pooling_type))
super(ResNetIBN, self).__init__()
self.depth = depth
self.pretrained = pretrained
self.cut_at_pooling = cut_at_pooling
resnet = ResNetIBN.__factory[depth](pretrained=pretrained)
resnet.layer4[0].conv2.stride = (1, 1)
resnet.layer4[0].downsample[0].stride = (1, 1)
self.base = nn.Sequential(
resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4)
self.gap = build_pooling_layer(pooling_type)
if not self.cut_at_pooling:
self.num_features = num_features
self.norm = norm
self.dropout = dropout
self.has_embedding = num_features > 0
self.num_classes = num_classes
out_planes = resnet.fc.in_features
# Append new layers
if self.has_embedding:
self.feat = nn.Linear(out_planes, self.num_features)
self.feat_bn = nn.BatchNorm1d(self.num_features)
init.kaiming_normal_(self.feat.weight, mode='fan_out')
init.constant_(self.feat.bias, 0)
else:
# Change the num_features to CNN output channels
self.num_features = out_planes
self.feat_bn = nn.BatchNorm1d(self.num_features)
self.feat_bn.bias.requires_grad_(False)
if self.dropout > 0:
self.drop = nn.Dropout(self.dropout)
if self.num_classes > 0:
self.classifier = nn.Linear(self.num_features, self.num_classes, bias=False)
init.normal_(self.classifier.weight, std=0.001)
init.constant_(self.feat_bn.weight, 1)
init.constant_(self.feat_bn.bias, 0)
if not pretrained:
self.reset_params()
def forward(self, x):
x = self.base(x)
x = self.gap(x)
x = x.view(x.size(0), -1)
if self.cut_at_pooling:
return x
if self.has_embedding:
bn_x = self.feat_bn(self.feat(x))
else:
bn_x = self.feat_bn(x)
if self.training is False:
bn_x = F.normalize(bn_x)
return bn_x
if self.norm:
bn_x = F.normalize(bn_x)
elif self.has_embedding:
bn_x = F.relu(bn_x)
if self.dropout > 0:
bn_x = self.drop(bn_x)
if self.num_classes > 0:
prob = self.classifier(bn_x)
else:
return bn_x
return prob, bn_x
def reset_params(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm2d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.BatchNorm1d):
init.constant_(m.weight, 1)
init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
init.normal_(m.weight, std=0.001)
if m.bias is not None:
init.constant_(m.bias, 0)
def resnet_ibn50a(**kwargs):
return ResNetIBN('50a', **kwargs)
def resnet_ibn101a(**kwargs):
return ResNetIBN('101a', **kwargs)
================================================
FILE: hhcl/models/resnet_ibn_a.py
================================================
import torch
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
__all__ = ['ResNet', 'resnet50_ibn_a', 'resnet101_ibn_a']
model_urls = {
'ibn_resnet50a': './examples/pretrained/resnet50_ibn_a.pth.tar',
'ibn_resnet101a': './examples/pretrained/resnet101_ibn_a.pth.tar',
}
def conv3x3(in_planes, out_planes, stride=1):
"3x3 convolution with padding"
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class IBN(nn.Module):
def __init__(self, planes):
super(IBN, self).__init__()
half1 = int(planes/2)
self.half = half1
half2 = planes - half1
self.IN = nn.InstanceNorm2d(half1, affine=True)
self.BN = nn.BatchNorm2d(half2)
def forward(self, x):
split = torch.split(x, self.half, 1)
out1 = self.IN(split[0].contiguous())
out2 = self.BN(split[1].contiguous())
out = torch.cat((out1, out2), 1)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
if ibn:
self.bn1 = IBN(planes)
else:
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, layers, num_classes=1000):
scale = 64
self.inplanes = scale
super(ResNet, self).__init__()
self.conv1 = nn.Conv2d(3, scale, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(scale)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, scale, layers[0])
self.layer2 = self._make_layer(block, scale*2, layers[1], stride=2)
self.layer3 = self._make_layer(block, scale*4, layers[2], stride=2)
self.layer4 = self._make_layer(block, scale*8, layers[3], stride=2)
self.avgpool = nn.AvgPool2d(7)
self.fc = nn.Linear(scale * 8 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.InstanceNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(self.inplanes, planes * block.expansion,
kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
ibn = True
if planes == 512:
ibn = False
layers.append(block(self.inplanes, planes, ibn, stride, downsample))
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(block(self.inplanes, planes, ibn))
return nn.Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
def resnet50_ibn_a(pretrained=False, **kwargs):
"""Constructs a ResNet-50 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
if pretrained:
state_dict = torch.load(model_urls['ibn_resnet50a'], map_location=torch.device('cpu'))['state_dict']
state_dict = remove_module_key(state_dict)
model.load_state_dict(state_dict)
return model
def resnet101_ibn_a(pretrained=False, **kwargs):
"""Constructs a ResNet-101 model.
Args:
pretrained (bool): If True, returns a model pre-trained on ImageNet
"""
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
if pretrained:
state_dict = torch.load(model_urls['ibn_resnet101a'], map_location=torch.device('cpu'))['state_dict']
state_dict = remove_module_key(state_dict)
model.load_state_dict(state_dict)
return model
def remove_module_key(state_dict):
for key in list(state_dict.keys()):
if 'module' in key:
state_dict[key.replace('module.','')] = state_dict.pop(key)
return state_dict
================================================
FILE: hhcl/models/triplet.py
================================================
from __future__ import absolute_import
import torch
from torch import nn
import torch.nn.functional as F
def euclidean_dist(x, y):
m, n = x.size(0), y.size(0)
xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
dist = xx + yy
dist.addmm_(1, -2, x, y.t())
dist = dist.clamp(min=1e-12).sqrt() # for numerical stability
return dist
def cosine_dist(x, y):
bs1, bs2 = x.size(0), y.size(0)
frac_up = torch.matmul(x, y.transpose(0, 1))
frac_down = (torch.sqrt(torch.sum(torch.pow(x, 2), 1))).view(bs1, 1).repeat(1, bs2) * \
(torch.sqrt(torch.sum(torch.pow(y, 2), 1))).view(1, bs2).repeat(bs1, 1)
cosine = frac_up / frac_down
return 1-cosine
def _batch_hard(mat_distance, mat_similarity, indice=False):
sorted_mat_distance, positive_indices = torch.sort(mat_distance + (-9999999.) * (1 - mat_similarity), dim=1, descending=True)
hard_p = sorted_mat_distance[:, 0]
hard_p_indice = positive_indices[:, 0]
sorted_mat_distance, negative_indices = torch.sort(mat_distance + (9999999.) * (mat_similarity), dim=1, descending=False)
hard_n = sorted_mat_distance[:, 0]
hard_n_indice = negative_indices[:, 0]
if(indice):
return hard_p, hard_n, hard_p_indice, hard_n_indice
return hard_p, hard_n
class TripletLoss(nn.Module):
'''
Compute Triplet loss augmented with Batch Hard
Details can be seen in 'In defense of the Triplet Loss for Person Re-Identification'
'''
def __init__(self, margin, normalize_feature=False):
super(TripletLoss, self).__init__()
self.margin = margin
self.normalize_feature = normalize_feature
self.margin_loss = nn.MarginRankingLoss(margin=margin).cuda()
def forward(self, emb, label):
if self.normalize_feature:
# equal to cosine similarity
emb = F.normalize(emb)
mat_dist = euclidean_dist(emb, emb)
# mat_dist = cosine_dist(emb, emb)
assert mat_dist.size(0) == mat_dist.size(1)
N = mat_dist.size(0)
mat_sim = label.expand(N, N).eq(label.expand(N, N).t()).float()
dist_ap, dist_an = _batch_hard(mat_dist, mat_sim)
assert dist_an.size(0)==dist_ap.size(0)
y = torch.ones_like(dist_ap)
loss = self.margin_loss(dist_an, dist_ap, y)
prec = (dist_an.data > dist_ap.data).sum() * 1. / y.size(0)
return loss, prec
class SoftTripletLoss(nn.Module):
def __init__(self, margin=None, normalize_feature=False):
super(SoftTripletLoss, self).__init__()
self.margin = margin
self.normalize_feature = normalize_feature
def forward(self, emb1, emb2, label):
if self.normalize_feature:
# equal to cosine similarity
emb1 = F.normalize(emb1)
emb2 = F.normalize(emb2)
mat_dist = euclidean_dist(emb1, emb1)
assert mat_dist.size(0) == mat_dist.size(1)
N = mat_dist.size(0)
mat_sim = label.expand(N, N).eq(label.expand(N, N).t()).float()
dist_ap, dist_an, ap_idx, an_idx = _batch_hard(mat_dist, mat_sim, indice=True)
assert dist_an.size(0)==dist_ap.size(0)
triple_dist = torch.stack((dist_ap, dist_an), dim=1)
triple_dist = F.log_softmax(triple_dist, dim=1)
if (self.margin is not None):
loss = (- self.margin * triple_dist[:,0] - (1 - self.margin) * triple_dist[:,1]).mean()
return loss
mat_dist_ref = euclidean_dist(emb2, emb2)
dist_ap_ref = torch.gather(mat_dist_ref, 1, ap_idx.view(N,1).expand(N,N))[:,0]
dist_an_ref = torch.gather(mat_dist_ref, 1, an_idx.view(N,1).expand(N,N))[:,0]
triple_dist_ref = torch.stack((dist_ap_ref, dist_an_ref), dim=1)
triple_dist_ref = F.softmax(triple_dist_ref, dim=1).detach()
loss = (- triple_dist_ref * triple_dist).mean(0).sum()
return loss
================================================
FILE: hhcl/trainers.py
================================================
from __future__ import print_function, absolute_import
import time
import torch
import torch.nn.functional as F
from .utils.meters import AverageMeter
class Trainer(object):
def __init__(self, encoder, memory=None):
super(Trainer, self).__init__()
self.encoder = encoder
self.memory = memory
def train(self, epoch, data_loader, optimizer, print_freq=10, train_iters=400):
self.encoder.train()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
end = time.time()
for i in range(train_iters):
# load data
inputs = data_loader.next()
data_time.update(time.time() - end)
# process inputs
inputs, labels, indexes = self._parse_data(inputs)
loss = 0
# forward
f_out = self._forward(inputs)
loss += self.memory(f_out, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.update(loss.item())
# print log
batch_time.update(time.time() - end)
end = time.time()
if (i + 1) % print_freq == 0:
print('Epoch: [{}][{}/{}]\t'
'Time {:.3f} ({:.3f})\t'
'Data {:.3f} ({:.3f})\t'
'Loss {:.3f} ({:.3f})'
.format(epoch, i + 1, len(data_loader),
batch_time.val, batch_time.avg,
data_time.val, data_time.avg,
losses.val, losses.avg))
def _parse_data(self, inputs):
imgs, _, pids, _, indexes = inputs
return imgs.cuda(), pids.cuda(), indexes.cuda()
def _forward(self, inputs):
return self.encoder(inputs)
================================================
FILE: hhcl/utils/__init__.py
================================================
from __future__ import absolute_import
import torch
def to_numpy(tensor):
if torch.is_tensor(tensor):
return tensor.cpu().numpy()
elif type(tensor).__module__ != 'numpy':
raise ValueError("Cannot convert {} to numpy array"
.format(type(tensor)))
return tensor
def to_torch(ndarray):
if type(ndarray).__module__ == 'numpy':
return torch.from_numpy(ndarray)
elif not torch.is_tensor(ndarray):
raise ValueError("Cannot convert {} to torch tensor"
.format(type(ndarray)))
return ndarray
================================================
FILE: hhcl/utils/data/__init__.py
================================================
from __future__ import absolute_import
from .base_dataset import BaseDataset, BaseImageDataset
from .preprocessor import Preprocessor
class IterLoader:
def __init__(self, loader, length=None):
self.loader = loader
self.length = length
self.iter = None
def __len__(self):
if self.length is not None:
return self.length
return len(self.loader)
def new_epoch(self):
self.iter = iter(self.loader)
def next(self):
try:
return next(self.iter)
except:
self.iter = iter(self.loader)
return next(self.iter)
================================================
FILE: hhcl/utils/data/base_dataset.py
================================================
# encoding: utf-8
import numpy as np
class BaseDataset(object):
"""
Base class of reid dataset
"""
def get_imagedata_info(self, data):
pids, cams = [], []
for _, pid, camid in data:
pids += [pid]
cams += [camid]
pids = set(pids)
cams = set(cams)
num_pids = len(pids)
num_cams = len(cams)
num_imgs = len(data)
return num_pids, num_imgs, num_cams
def print_dataset_statistics(self):
raise NotImplementedError
@property
def images_dir(self):
return None
class BaseImageDataset(BaseDataset):
"""
Base class of image reid dataset
"""
def print_dataset_statistics(self, train, query, gallery):
num_train_pids, num_train_imgs, num_train_cams = self.get_imagedata_info(train)
num_query_pids, num_query_imgs, num_query_cams = self.get_imagedata_info(query)
num_gallery_pids, num_gallery_imgs, num_gallery_cams = self.get_imagedata_info(gallery)
print("Dataset statistics:")
print(" ----------------------------------------")
print(" subset | # ids | # images | # cameras")
print(" ----------------------------------------")
print(" train | {:5d} | {:8d} | {:9d}".format(num_train_pids, num_train_imgs, num_train_cams))
print(" query | {:5d} | {:8d} | {:9d}".format(num_query_pids, num_query_imgs, num_query_cams))
print(" gallery | {:5d} | {:8d} | {:9d}".format(num_gallery_pids, num_gallery_imgs, num_gallery_cams))
print(" ----------------------------------------")
================================================
FILE: hhcl/utils/data/preprocessor.py
================================================
from __future__ import absolute_import
import os
import os.path as osp
from torch.utils.data import DataLoader, Dataset
import numpy as np
import random
import math
from PIL import Image
class Preprocessor(Dataset):
def __init__(self, dataset, root=None, transform=None, mutual=False):
super(Preprocessor, self).__init__()
self.dataset = dataset
self.root = root
self.transform = transform
self.mutual = mutual
def __len__(self):
return len(self.dataset)
def __getitem__(self, indices):
if self.mutual:
return self._get_mutual_item(indices)
else:
return self._get_single_item(indices)
def _get_single_item(self, index):
fname, pid, camid = self.dataset[index]
fpath = fname
if self.root is not None:
fpath = osp.join(self.root, fname)
img = Image.open(fpath).convert('RGB')
if self.transform is not None:
img = self.transform(img)
return img, fname, pid, camid, index
def _get_mutual_item(self, index):
fname, pid, camid = self.dataset[index]
fpath = fname
if self.root is not None:
fpath = osp.join(self.root, fname)
img_1 = Image.open(fpath).convert('RGB')
img_2 = img_1.copy()
if self.transform is not None:
img_1 = self.transform(img_1)
img_2 = self.transform(img_2)
return img_1, img_2, pid, camid
================================================
FILE: hhcl/utils/data/sampler.py
================================================
from __future__ import absolute_import
from collections import defaultdict
import math
import numpy as np
import copy
import random
import torch
from torch.utils.data.sampler import (
Sampler, SequentialSampler, RandomSampler, SubsetRandomSampler,
WeightedRandomSampler)
def No_index(a, b):
assert isinstance(a, list)
return [i for i, j in enumerate(a) if j != b]
class RandomIdentitySampler(Sampler):
def __init__(self, data_source, num_instances):
self.data_source = data_source
self.num_instances = num_instances
self.index_dic = defaultdict(list)
for index, (_, pid, _) in enumerate(data_source):
self.index_dic[pid].append(index)
self.pids = list(self.index_dic.keys())
self.num_samples = len(self.pids)
def __len__(self):
return self.num_samples * self.num_instances
def __iter__(self):
indices = torch.randperm(self.num_samples).tolist()
ret = []
for i in indices:
pid = self.pids[i]
t = self.index_dic[pid]
if len(t) >= self.num_instances:
t = np.random.choice(t, size=self.num_instances, replace=False)
else:
t = np.random.choice(t, size=self.num_instances, replace=True)
ret.extend(t)
return iter(ret)
class RandomMultipleGallerySampler(Sampler):
def __init__(self, data_source, num_instances=4):
super().__init__(data_source)
self.data_source = data_source
self.index_pid = defaultdict(int)
self.pid_cam = defaultdict(list)
self.pid_index = defaultdict(list)
self.num_instances = num_instances
for index, (_, pid, cam) in enumerate(data_source):
if pid < 0:
continue
self.index_pid[index] = pid
self.pid_cam[pid].append(cam)
self.pid_index[pid].append(index)
self.pids = list(self.pid_index.keys())
self.num_samples = len(self.pids)
def __len__(self):
return self.num_samples * self.num_instances
def __iter__(self):
indices = torch.randperm(len(self.pids)).tolist()
ret = []
for kid in indices:
i = random.choice(self.pid_index[self.pids[kid]])
_, i_pid, i_cam = self.data_source[i]
ret.append(i)
pid_i = self.index_pid[i]
cams = self.pid_cam[pid_i]
index = self.pid_index[pid_i]
select_cams = No_index(cams, i_cam)
if select_cams:
if len(select_cams) >= self.num_instances:
cam_indexes = np.random.choice(select_cams, size=self.num_instances-1, replace=False)
else:
cam_indexes = np.random.choice(select_cams, size=self.num_instances-1, replace=True)
for kk in cam_indexes:
ret.append(index[kk])
else:
select_indexes = No_index(index, i)
if not select_indexes:
continue
if len(select_indexes) >= self.num_instances:
ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=False)
else:
ind_indexes = np.random.choice(select_indexes, size=self.num_instances-1, replace=True)
for kk in ind_indexes:
ret.append(index[kk])
return iter(ret)
================================================
FILE: hhcl/utils/data/transforms.py
================================================
from __future__ import absolute_import
from torchvision.transforms import *
from PIL import Image
import random
import math
import numpy as np
class RectScale(object):
def __init__(self, height, width, interpolation=Image.BILINEAR):
self.height = height
self.width = width
self.interpolation = interpolation
def __call__(self, img):
w, h = img.size
if h == self.height and w == self.width:
return img
return img.resize((self.width, self.height), self.interpolation)
class RandomSizedRectCrop(object):
def __init__(self, height, width, interpolation=Image.BILINEAR):
self.height = height
self.width = width
self.interpolation = interpolation
def __call__(self, img):
for attempt in range(10):
area = img.size[0] * img.size[1]
target_area = random.uniform(0.64, 1.0) * area
aspect_ratio = random.uniform(2, 3)
h = int(round(math.sqrt(target_area * aspect_ratio)))
w = int(round(math.sqrt(target_area / aspect_ratio)))
if w <= img.size[0] and h <= img.size[1]:
x1 = random.randint(0, img.size[0] - w)
y1 = random.randint(0, img.size[1] - h)
img = img.crop((x1, y1, x1 + w, y1 + h))
assert(img.size == (w, h))
return img.resize((self.width, self.height), self.interpolation)
# Fallback
scale = RectScale(self.height, self.width,
interpolation=self.interpolation)
return scale(img)
class RandomErasing(object):
""" Randomly selects a rectangle region in an image and erases its pixels.
'Random Erasing Data Augmentation' by Zhong et al.
See https://arxiv.org/pdf/1708.04896.pdf
Args:
probability: The probability that the Random Erasing operation will be performed.
sl: Minimum proportion of erased area against input image.
sh: Maximum proportion of erased area against input image.
r1: Minimum aspect ratio of erased area.
mean: Erasing value.
"""
def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4914, 0.4822, 0.4465)):
self.probability = probability
self.mean = mean
self.sl = sl
self.sh = sh
self.r1 = r1
def __call__(self, img):
if random.uniform(0, 1) >= self.probability:
return img
for attempt in range(100):
area = img.size()[1] * img.size()[2]
target_area = random.uniform(self.sl, self.sh) * area
aspect_ratio = random.uniform(self.r1, 1 / self.r1)
h = int(round(math.sqrt(target_area * aspect_ratio)))
w = int(round(math.sqrt(target_area / aspect_ratio)))
if w < img.size()[2] and h < img.size()[1]:
x1 = random.randint(0, img.size()[1] - h)
y1 = random.randint(0, img.size()[2] - w)
if img.size()[0] == 3:
img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
img[1, x1:x1 + h, y1:y1 + w] = self.mean[1]
img[2, x1:x1 + h, y1:y1 + w] = self.mean[2]
else:
img[0, x1:x1 + h, y1:y1 + w] = self.mean[0]
return img
return img
================================================
FILE: hhcl/utils/faiss_rerank.py
================================================
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017.
url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf
Matlab version: https://github.com/zhunzhong07/person-re-ranking
"""
import os, sys
import time
import numpy as np
from scipy.spatial.distance import cdist
import gc
import faiss
import torch
import torch.nn.functional as F
from .faiss_utils import search_index_pytorch, search_raw_array_pytorch, \
index_init_gpu, index_init_cpu
def k_reciprocal_neigh(initial_rank, i, k1):
forward_k_neigh_index = initial_rank[i,:k1+1]
backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1]
fi = np.where(backward_k_neigh_index==i)[0]
return forward_k_neigh_index[fi]
def compute_jaccard_distance(target_features, k1=20, k2=6, print_flag=True, search_option=0, use_float16=False):
end = time.time()
if print_flag:
print('Computing jaccard distance...')
ngpus = faiss.get_num_gpus()
N = target_features.size(0)
mat_type = np.float16 if use_float16 else np.float32
if (search_option==0):
# GPU + PyTorch CUDA Tensors (1)
res = faiss.StandardGpuResources()
res.setDefaultNullStreamAllDevices()
_, initial_rank = search_raw_array_pytorch(res, target_features, target_features, k1)
initial_rank = initial_rank.cpu().numpy()
elif (search_option==1):
# GPU + PyTorch CUDA Tensors (2)
res = faiss.StandardGpuResources()
index = faiss.GpuIndexFlatL2(res, target_features.size(-1))
index.add(target_features.cpu().numpy())
_, initial_rank = search_index_pytorch(index, target_features, k1)
res.syncDefaultStreamCurrentDevice()
initial_rank = initial_rank.cpu().numpy()
elif (search_option==2):
# GPU
index = index_init_gpu(ngpus, target_features.size(-1))
index.add(target_features.cpu().numpy())
_, initial_rank = index.search(target_features.cpu().numpy(), k1)
else:
# CPU
index = index_init_cpu(target_features.size(-1))
index.add(target_features.cpu().numpy())
_, initial_rank = index.search(target_features.cpu().numpy(), k1)
nn_k1 = []
nn_k1_half = []
for i in range(N):
nn_k1.append(k_reciprocal_neigh(initial_rank, i, k1))
nn_k1_half.append(k_reciprocal_neigh(initial_rank, i, int(np.around(k1/2))))
V = np.zeros((N, N), dtype=mat_type)
for i in range(N):
k_reciprocal_index = nn_k1[i]
k_reciprocal_expansion_index = k_reciprocal_index
for candidate in k_reciprocal_index:
candidate_k_reciprocal_index = nn_k1_half[candidate]
if (len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index)) > 2/3*len(candidate_k_reciprocal_index)):
k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index)
k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) ## element-wise unique
dist = 2-2*torch.mm(target_features[i].unsqueeze(0).contiguous(), target_features[k_reciprocal_expansion_index].t())
if use_float16:
V[i,k_reciprocal_expansion_index] = F.softmax(-dist, dim=1).view(-1).cpu().numpy().astype(mat_type)
else:
V[i,k_reciprocal_expansion_index] = F.softmax(-dist, dim=1).view(-1).cpu().numpy()
del nn_k1, nn_k1_half
if k2 != 1:
V_qe = np.zeros_like(V, dtype=mat_type)
for i in range(N):
V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:], axis=0)
V = V_qe
del V_qe
del initial_rank
invIndex = []
for i in range(N):
invIndex.append(np.where(V[:,i] != 0)[0]) #len(invIndex)=all_num
jaccard_dist = np.zeros((N, N), dtype=mat_type)
for i in range(N):
temp_min = np.zeros((1, N), dtype=mat_type)
# temp_max = np.zeros((1,N), dtype=mat_type)
indNonZero = np.where(V[i, :] != 0)[0]
indImages = []
indImages = [invIndex[ind] for ind in indNonZero]
for j in range(len(indNonZero)):
temp_min[0, indImages[j]] = temp_min[0, indImages[j]]+np.minimum(V[i, indNonZero[j]], V[indImages[j], indNonZero[j]])
# temp_max[0,indImages[j]] = temp_max[0,indImages[j]]+np.maximum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]])
jaccard_dist[i] = 1-temp_min/(2-temp_min)
# jaccard_dist[i] = 1-temp_min/(temp_max+1e-6)
del invIndex, V
pos_bool = (jaccard_dist < 0)
jaccard_dist[pos_bool] = 0.0
if print_flag:
print("Jaccard distance computing time cost: {}".format(time.time()-end))
return jaccard_dist
================================================
FILE: hhcl/utils/faiss_utils.py
================================================
import os
import numpy as np
import faiss
import torch
def swig_ptr_from_FloatTensor(x):
assert x.is_contiguous()
assert x.dtype == torch.float32
return faiss.cast_integer_to_float_ptr(
x.storage().data_ptr() + x.storage_offset() * 4)
def swig_ptr_from_LongTensor(x):
assert x.is_contiguous()
assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
return faiss.cast_integer_to_long_ptr(
x.storage().data_ptr() + x.storage_offset() * 8)
def search_index_pytorch(index, x, k, D=None, I=None):
"""call the search function of an index with pytorch tensor I/O (CPU
and GPU supported)"""
assert x.is_contiguous()
n, d = x.size()
assert d == index.d
if D is None:
D = torch.empty((n, k), dtype=torch.float32, device=x.device)
else:
assert D.size() == (n, k)
if I is None:
I = torch.empty((n, k), dtype=torch.int64, device=x.device)
else:
assert I.size() == (n, k)
torch.cuda.synchronize()
xptr = swig_ptr_from_FloatTensor(x)
Iptr = swig_ptr_from_LongTensor(I)
Dptr = swig_ptr_from_FloatTensor(D)
index.search_c(n, xptr,
k, Dptr, Iptr)
torch.cuda.synchronize()
return D, I
def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,
metric=faiss.METRIC_L2):
assert xb.device == xq.device
nq, d = xq.size()
if xq.is_contiguous():
xq_row_major = True
elif xq.t().is_contiguous():
xq = xq.t() # I initially wrote xq:t(), Lua is still haunting me :-)
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xq_ptr = swig_ptr_from_FloatTensor(xq)
nb, d2 = xb.size()
assert d2 == d
if xb.is_contiguous():
xb_row_major = True
elif xb.t().is_contiguous():
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
xb_ptr = swig_ptr_from_FloatTensor(xb)
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
assert D.device == xb.device
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
else:
assert I.shape == (nq, k)
assert I.device == xb.device
D_ptr = swig_ptr_from_FloatTensor(D)
I_ptr = swig_ptr_from_LongTensor(I)
faiss.bruteForceKnn(res, metric,
xb_ptr, xb_row_major, nb,
xq_ptr, xq_row_major, nq,
d, k, D_ptr, I_ptr)
return D, I
def index_init_gpu(ngpus, feat_dim):
flat_config = []
for i in range(ngpus):
cfg = faiss.GpuIndexFlatConfig()
cfg.useFloat16 = False
cfg.device = i
flat_config.append(cfg)
res = [faiss.StandardGpuResources() for i in range(ngpus)]
indexes = [faiss.GpuIndexFlatL2(res[i], feat_dim, flat_config[i]) for i in range(ngpus)]
index = faiss.IndexShards(feat_dim)
for sub_index in indexes:
index.add_shard(sub_index)
index.reset()
return index
def index_init_cpu(feat_dim):
return faiss.IndexFlatL2(feat_dim)
================================================
FILE: hhcl/utils/logging.py
================================================
from __future__ import absolute_import
import os
import sys
from .osutils import mkdir_if_missing
class Logger(object):
def __init__(self, fpath=None):
self.console = sys.stdout
self.file = None
if fpath is not None:
mkdir_if_missing(os.path.dirname(fpath))
self.file = open(fpath, 'w')
def __del__(self):
self.close()
def __enter__(self):
pass
def __exit__(self, *args):
self.close()
def write(self, msg):
self.console.write(msg)
if self.file is not None:
self.file.write(msg)
def flush(self):
self.console.flush()
if self.file is not None:
self.file.flush()
os.fsync(self.file.fileno())
def close(self):
self.console.close()
if self.file is not None:
self.file.close()
================================================
FILE: hhcl/utils/meters.py
================================================
from __future__ import absolute_import
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
================================================
FILE: hhcl/utils/osutils.py
================================================
from __future__ import absolute_import
import os
import errno
def mkdir_if_missing(dir_path):
try:
os.makedirs(dir_path)
except OSError as e:
if e.errno != errno.EEXIST:
raise
================================================
FILE: hhcl/utils/rerank.py
================================================
#!/usr/bin/env python2/python3
# -*- coding: utf-8 -*-
"""
Source: https://github.com/zhunzhong07/person-re-ranking
Created on Mon Jun 26 14:46:56 2017
@author: luohao
Modified by Houjing Huang, 2017-12-22.
- This version accepts distance matrix instead of raw features.
- The difference of `/` division between python 2 and 3 is handled.
- numpy.float16 is replaced by numpy.float32 for numerical precision.
CVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017.
url:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf
Matlab version: https://github.com/zhunzhong07/person-re-ranking
API
q_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery]
q_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query]
g_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery]
k1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3)
Returns:
final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery]
"""
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
__all__ = ['re_ranking']
import numpy as np
def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3):
# The following naming, e.g. gallery_num, is different from outer scope.
# Don't care about it.
original_dist = np.concatenate(
[np.concatenate([q_q_dist, q_g_dist], axis=1),
np.concatenate([q_g_dist.T, g_g_dist], axis=1)],
axis=0)
original_dist = np.power(original_dist, 2).astype(np.float32)
original_dist = np.transpose(1. * original_dist/np.max(original_dist,axis = 0))
V = np.zeros_like(original_dist).astype(np.float32)
initial_rank = np.argsort(original_dist).astype(np.int32)
query_num = q_g_dist.shape[0]
gallery_num = q_g_dist.shape[0] + q_g_dist.shape[1]
all_num = gallery_num
for i in range(all_num):
# k-reciprocal neighbors
forward_k_neigh_index = initial_rank[i,:k1+1]
backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1]
fi = np.where(backward_k_neigh_index==i)[0]
k_reciprocal_index = forward_k_neigh_index[fi]
k_reciprocal_expansion_index = k_reciprocal_index
for j in range(len(k_reciprocal_index)):
candidate = k_reciprocal_index[j]
candidate_forward_k_neigh_index = initial_rank[candidate,:int(np.around(k1/2.))+1]
candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,:int(np.around(k1/2.))+1]
fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0]
candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate]
if len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index))> 2./3*len(candidate_k_reciprocal_index):
k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index)
k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index)
weight = np.exp(-original_dist[i,k_reciprocal_expansion_index])
V[i,k_reciprocal_expansion_index] = 1.*weight/np.sum(weight)
original_dist = original_dist[:query_num,]
if k2 != 1:
V_qe = np.zeros_like(V,dtype=np.float32)
for i in range(all_num):
V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:],axis=0)
V = V_qe
del V_qe
del initial_rank
invIndex = []
for i in range(gallery_num):
invIndex.append(np.where(V[:,i] != 0)[0])
jaccard_dist = np.zeros_like(original_dist,dtype = np.float32)
for i in range(query_num):
temp_min = np.zeros(shape=[1,gallery_num],dtype=np.float32)
indNonZero = np.where(V[i,:] != 0)[0]
indImages = []
indImages = [invIndex[ind] for ind in indNonZero]
for j in range(len(indNonZero)):
temp_min[0,indImages[j]] = temp_min[0,indImages[j]]+ np.minimum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]])
jaccard_dist[i] = 1-temp_min/(2.-temp_min)
final_dist = jaccard_dist*(1-lambda_value) + original_dist*lambda_value
del original_dist
del V
del jaccard_dist
final_dist = final_dist[:query_num,query_num:]
return final_dist
================================================
FILE: hhcl/utils/serialization.py
================================================
from __future__ import print_function, absolute_import
import json
import os.path as osp
import shutil
import torch
from torch.nn import Parameter
from .osutils import mkdir_if_missing
def read_json(fpath):
with open(fpath, 'r') as f:
obj = json.load(f)
return obj
def write_json(obj, fpath):
mkdir_if_missing(osp.dirname(fpath))
with open(fpath, 'w') as f:
json.dump(obj, f, indent=4, separators=(',', ': '))
def save_checkpoint(state, is_best, fpath='checkpoint.pth.tar'):
mkdir_if_missing(osp.dirname(fpath))
torch.save(state, fpath)
if is_best:
shutil.copy(fpath, osp.join(osp.dirname(fpath), 'model_best.pth.tar'))
def load_checkpoint(fpath):
if osp.isfile(fpath):
# checkpoint = torch.load(fpath)
checkpoint = torch.load(fpath, map_location=torch.device('cpu'))
print("=> Loaded checkpoint '{}'".format(fpath))
return checkpoint
else:
raise ValueError("=> No checkpoint found at '{}'".format(fpath))
def copy_state_dict(state_dict, model, strip=None):
tgt_state = model.state_dict()
copied_names = set()
for name, param in state_dict.items():
if strip is not None and name.startswith(strip):
name = name[len(strip):]
if name not in tgt_state:
continue
if isinstance(param, Parameter):
param = param.data
if param.size() != tgt_state[name].size():
print('mismatch:', name, param.size(), tgt_state[name].size())
continue
tgt_state[name].copy_(param)
copied_names.add(name)
missing = set(tgt_state.keys()) - copied_names
if len(missing) > 0:
print("missing keys in state_dict:", missing)
return model
================================================
FILE: requirements.txt
================================================
numpy
sklearn
Cython
h5py
pyzmq
pillow-simd
six
scipy
matplotlib
faiss-gpu==1.6.3
easydict
================================================
FILE: run.sh
================================================
### resnet50 ###
# market1501
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d market1501 --iters 200 --eps 0.45 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_avg_cmhybrid
# dukemtmcreid
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d dukemtmcreid --iters 200 --eps 0.6 --num-instances 16 --pooling-type avg --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_avg_cmhybrid
### resnet_ibn50a + gem pooling ###
# market1501
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet_ibn50a -d market1501 --iters 200 --eps 0.45 --num-instances 16 --pooling-type gem --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/market1501/resnet50_ibn_gem_cmhybrid
# dukemtmcreid
CUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet_ibn50a -d dukemtmcreid --iters 200 --eps 0.6 --num-instances 16 --pooling-type gem --memorybank CMhybrid --epochs 60 --logs-dir examples/logs/dukemtmcreid/resnet50_ibn_gem_cmhybrid
# test
CUDA_VISIBLE_DEVICES=0 python examples/test.py -d market1501 --data-dir examples/data/market1501 --pooling-type avg --resume examples/logs/market1501/resnet50_avg_cmhybrid/model_best.pth.tar
================================================
FILE: setup.py
================================================
from setuptools import setup, find_packages
setup(name='hhcl',
version='1.0.0',
install_requires=[
'numpy', 'torch', 'torchvision',
'six', 'h5py', 'Pillow', 'scipy',
'scikit-learn', 'metric-learn', 'faiss_gpu'],
packages=find_packages(),
keywords=[
'Unsupervised Learning',
'Contrastive Learning',
'Object Re-identification'
])
gitextract_d7xkgx8j/ ├── README.md ├── examples/ │ ├── test.py │ └── train.py ├── hhcl/ │ ├── __init__.py │ ├── datasets/ │ │ ├── __init__.py │ │ ├── celebreid.py │ │ ├── dukemtmcreid.py │ │ ├── market1501.py │ │ ├── msmt17.py │ │ └── personx.py │ ├── evaluation_metrics/ │ │ ├── __init__.py │ │ ├── classification.py │ │ └── ranking.py │ ├── evaluators.py │ ├── models/ │ │ ├── __init__.py │ │ ├── cm.py │ │ ├── dsbn.py │ │ ├── kmeans.py │ │ ├── losses.py │ │ ├── pooling.py │ │ ├── resnet.py │ │ ├── resnet_ibn.py │ │ ├── resnet_ibn_a.py │ │ └── triplet.py │ ├── trainers.py │ └── utils/ │ ├── __init__.py │ ├── data/ │ │ ├── __init__.py │ │ ├── base_dataset.py │ │ ├── preprocessor.py │ │ ├── sampler.py │ │ └── transforms.py │ ├── faiss_rerank.py │ ├── faiss_utils.py │ ├── logging.py │ ├── meters.py │ ├── osutils.py │ ├── rerank.py │ └── serialization.py ├── requirements.txt ├── run.sh └── setup.py
SYMBOL INDEX (248 symbols across 35 files)
FILE: examples/test.py
function get_data (line 23) | def get_data(name, data_dir, height, width, batch_size, workers):
function main (line 45) | def main():
function main_worker (line 57) | def main_worker(args):
FILE: examples/train.py
function get_data (line 36) | def get_data(name, data_dir):
function get_train_loader (line 42) | def get_train_loader(args, dataset, height, width, batch_size, workers,
function get_test_loader (line 72) | def get_test_loader(dataset, height, width, batch_size, workers, testset...
function create_model (line 93) | def create_model(args):
function main (line 110) | def main():
function main_worker (line 122) | def main_worker(args):
FILE: hhcl/datasets/__init__.py
function names (line 20) | def names():
function create (line 24) | def create(name, root, *args, **kwargs):
function get_dataset (line 47) | def get_dataset(name, root, *args, **kwargs):
FILE: hhcl/datasets/celebreid.py
class CelebReID (line 8) | class CelebReID(BaseImageDataset):
method __init__ (line 14) | def __init__(self, root, verbose=True, **kwargs):
method _check_before_run (line 39) | def _check_before_run(self):
method _process_dir (line 50) | def _process_dir(self, dir_path, relabel=False):
FILE: hhcl/datasets/dukemtmcreid.py
function process_dir (line 7) | def process_dir(dir_path, relabel=False):
class DukeMTMCreID (line 37) | class DukeMTMCreID(BaseImageDataset):
method __init__ (line 55) | def __init__(self, root, verbose=True):
method _check_before_run (line 76) | def _check_before_run(self):
FILE: hhcl/datasets/market1501.py
class Market1501 (line 8) | class Market1501(BaseImageDataset):
method __init__ (line 21) | def __init__(self, root, verbose=True, **kwargs):
method _check_before_run (line 46) | def _check_before_run(self):
method _process_dir (line 57) | def _process_dir(self, dir_path, relabel=False):
FILE: hhcl/datasets/msmt17.py
function _process_dir (line 9) | def _process_dir(dir_path, relabel=False):
class MSMT17 (line 34) | class MSMT17(BaseImageDataset):
method __init__ (line 37) | def __init__(self, root, verbose=True, **kwargs):
method _check_before_run (line 62) | def _check_before_run(self):
FILE: hhcl/datasets/personx.py
class PersonX (line 9) | class PersonX(BaseImageDataset):
method __init__ (line 21) | def __init__(self, root, verbose=True, **kwargs):
method _check_before_run (line 46) | def _check_before_run(self):
method _process_dir (line 57) | def _process_dir(self, dir_path, relabel=False):
FILE: hhcl/evaluation_metrics/classification.py
function accuracy (line 7) | def accuracy(output, target, topk=(1,)):
FILE: hhcl/evaluation_metrics/ranking.py
function _unique_sample (line 10) | def _unique_sample(ids_dict, num):
function cmc (line 18) | def cmc(distmat, query_ids=None, gallery_ids=None,
function mean_ap (line 82) | def mean_ap(distmat, query_ids=None, gallery_ids=None,
FILE: hhcl/evaluators.py
function extract_cnn_feature (line 16) | def extract_cnn_feature(model, inputs):
function extract_features (line 23) | def extract_features(model, data_loader, print_freq=50):
function pairwise_distance (line 55) | def pairwise_distance(features, query=None, gallery=None):
function evaluate_all (line 75) | def evaluate_all(query_features, gallery_features, distmat, query=None, ...
class Evaluator (line 109) | class Evaluator(object):
method __init__ (line 110) | def __init__(self, model):
method evaluate (line 114) | def evaluate(self, data_loader, query, gallery, cmc_flag=False, rerank...
FILE: hhcl/models/__init__.py
function names (line 17) | def names():
function create (line 21) | def create(name, *args, **kwargs):
FILE: hhcl/models/cm.py
class CM (line 10) | class CM(autograd.Function):
method forward (line 13) | def forward(ctx, inputs, targets, features, momentum):
method backward (line 22) | def backward(ctx, grad_outputs):
function cm (line 36) | def cm(inputs, indexes, features, momentum=0.5):
class CM_Hard (line 40) | class CM_Hard(autograd.Function):
method forward (line 43) | def forward(ctx, inputs, targets, features, momentum):
method backward (line 52) | def backward(ctx, grad_outputs):
function cm_hard (line 75) | def cm_hard(inputs, indexes, features, momentum=0.5):
class CM_Hybrid (line 79) | class CM_Hybrid(autograd.Function):
method forward (line 82) | def forward(ctx, inputs, targets, features, momentum):
method backward (line 91) | def backward(ctx, grad_outputs):
function cm_hybrid (line 119) | def cm_hybrid(inputs, indexes, features, momentum=0.5):
class CM_Hybrid_v2 (line 123) | class CM_Hybrid_v2(autograd.Function):
method forward (line 126) | def forward(ctx, inputs, targets, features, momentum, num_instances):
method backward (line 136) | def backward(ctx, grad_outputs):
function cm_hybrid_v2 (line 162) | def cm_hybrid_v2(inputs, indexes, features, momentum=0.5, num_instances=...
class ClusterMemory (line 166) | class ClusterMemory(nn.Module, ABC):
method __init__ (line 173) | def __init__(self, num_features, num_samples, temp=0.05, momentum=0.2,...
method forward (line 201) | def forward(self, inputs, targets):
FILE: hhcl/models/dsbn.py
class DSBN2d (line 6) | class DSBN2d(nn.Module):
method __init__ (line 7) | def __init__(self, planes):
method forward (line 13) | def forward(self, x):
class DSBN1d (line 25) | class DSBN1d(nn.Module):
method __init__ (line 26) | def __init__(self, planes):
method forward (line 32) | def forward(self, x):
function convert_dsbn (line 44) | def convert_dsbn(model):
function convert_bn (line 60) | def convert_bn(model, use_target=True):
FILE: hhcl/models/kmeans.py
function label_generator_kmeans (line 14) | def label_generator_kmeans(features, num_classes=500, cuda=True):
FILE: hhcl/models/losses.py
class CrossEntropyLabelSmooth (line 7) | class CrossEntropyLabelSmooth(nn.Module):
method __init__ (line 17) | def __init__(self, num_classes=0, epsilon=0.1, topk_smoothing=False):
method forward (line 24) | def forward(self, inputs, targets):
class SoftEntropy (line 42) | class SoftEntropy(nn.Module):
method __init__ (line 43) | def __init__(self, input_prob=False):
method forward (line 48) | def forward(self, inputs, targets):
class SoftEntropySmooth (line 57) | class SoftEntropySmooth(nn.Module):
method __init__ (line 58) | def __init__(self, epsilon=0.1):
method forward (line 63) | def forward(self, inputs, soft_targets, targets):
class Softmax (line 72) | class Softmax(nn.Module):
method __init__ (line 74) | def __init__(self, feat_dim, num_class, temp=0.05):
method forward (line 80) | def forward(self, feats, labels):
class CircleLoss (line 89) | class CircleLoss(nn.Module):
method __init__ (line 93) | def __init__(self, feat_dim, num_class, margin=0.25, gamma=256):
method forward (line 105) | def forward(self, feats, labels):
class CosFace (line 131) | class CosFace(nn.Module):
method __init__ (line 140) | def __init__(self, feat_dim, num_class, s = 64.0, m = 0.35):
method forward (line 150) | def forward(self, input, label):
method __repr__ (line 165) | def __repr__(self):
class InstanceLoss (line 175) | class InstanceLoss(nn.Module):
method __init__ (line 176) | def __init__(self, batch_size, temperature, device):
method mask_correlated_samples (line 185) | def mask_correlated_samples(self, batch_size):
method forward (line 195) | def forward(self, z_i, z_j):
class ClusterLoss (line 214) | class ClusterLoss(nn.Module):
method __init__ (line 215) | def __init__(self, class_num, temperature, device):
method mask_correlated_clusters (line 225) | def mask_correlated_clusters(self, class_num):
method forward (line 235) | def forward(self, c_i, c_j):
class FocalLoss (line 264) | class FocalLoss(nn.Module):
method __init__ (line 265) | def __init__(self, gamma=2, alpha=0.25):
method forward (line 271) | def forward(self, input, target):
class LabelRefineLoss (line 288) | class LabelRefineLoss(nn.Module):
method __init__ (line 289) | def __init__(self, lambda1=0.0):
method forward (line 294) | def forward(self, input, target):
class FocalTopLoss (line 311) | class FocalTopLoss(nn.Module):
method __init__ (line 312) | def __init__(self, top_percent=0.7):
method masked_softmax_multi_focal (line 316) | def masked_softmax_multi_focal(self, vec, targets=None, dim=1):
method forward (line 345) | def forward(self, input, target):
FILE: hhcl/models/pooling.py
class GeneralizedMeanPoolingList (line 19) | class GeneralizedMeanPoolingList(nn.Module, ABC):
method __init__ (line 34) | def __init__(self, output_size=1, eps=1e-6):
method forward (line 39) | def forward(self, x_list):
method __repr__ (line 47) | def __repr__(self):
class GeneralizedMeanPooling (line 57) | class GeneralizedMeanPooling(nn.Module, ABC):
method __init__ (line 72) | def __init__(self, norm, output_size=1, eps=1e-6):
method forward (line 79) | def forward(self, x):
method __repr__ (line 85) | def __repr__(self):
class GeneralizedMeanPoolingP (line 97) | class GeneralizedMeanPoolingP(GeneralizedMeanPooling, ABC):
method __init__ (line 101) | def __init__(self, norm=3, output_size=1, eps=1e-6):
class GeneralizedMeanPoolingFpn (line 106) | class GeneralizedMeanPoolingFpn(nn.Module, ABC):
method __init__ (line 121) | def __init__(self, norm, output_size=1, eps=1e-6):
method forward (line 128) | def forward(self, x_lists):
method __repr__ (line 138) | def __repr__(self):
class GeneralizedMeanPoolingPFpn (line 150) | class GeneralizedMeanPoolingPFpn(GeneralizedMeanPoolingFpn, ABC):
method __init__ (line 154) | def __init__(self, norm=3, output_size=1, eps=1e-6):
class AdaptiveAvgMaxPool2d (line 159) | class AdaptiveAvgMaxPool2d(nn.Module, ABC):
method __init__ (line 160) | def __init__(self):
method forward (line 164) | def forward(self, x):
class FastGlobalAvgPool2d (line 171) | class FastGlobalAvgPool2d(nn.Module, ABC):
method __init__ (line 172) | def __init__(self, flatten=False):
method forward (line 176) | def forward(self, x):
function avg_pooling (line 188) | def avg_pooling():
function max_pooling (line 193) | def max_pooling():
class Flatten (line 197) | class Flatten(nn.Module):
method forward (line 198) | def forward(self, input):
function pooling_names (line 212) | def pooling_names():
function build_pooling_layer (line 216) | def build_pooling_layer(name):
FILE: hhcl/models/resnet.py
class ResNet (line 15) | class ResNet(nn.Module):
method __init__ (line 24) | def __init__(self, depth, pretrained=True, cut_at_pooling=False,
method forward (line 75) | def forward(self, x):
method reset_params (line 109) | def reset_params(self):
function resnet18 (line 127) | def resnet18(**kwargs):
function resnet34 (line 131) | def resnet34(**kwargs):
function resnet50 (line 135) | def resnet50(**kwargs):
function resnet101 (line 139) | def resnet101(**kwargs):
function resnet152 (line 143) | def resnet152(**kwargs):
FILE: hhcl/models/resnet_ibn.py
class ResNetIBN (line 16) | class ResNetIBN(nn.Module):
method __init__ (line 22) | def __init__(self, depth, pretrained=True, cut_at_pooling=False,
method forward (line 74) | def forward(self, x):
method reset_params (line 107) | def reset_params(self):
function resnet_ibn50a (line 125) | def resnet_ibn50a(**kwargs):
function resnet_ibn101a (line 129) | def resnet_ibn101a(**kwargs):
FILE: hhcl/models/resnet_ibn_a.py
function conv3x3 (line 16) | def conv3x3(in_planes, out_planes, stride=1):
class BasicBlock (line 22) | class BasicBlock(nn.Module):
method __init__ (line 25) | def __init__(self, inplanes, planes, stride=1, downsample=None):
method forward (line 35) | def forward(self, x):
class IBN (line 54) | class IBN(nn.Module):
method __init__ (line 55) | def __init__(self, planes):
method forward (line 63) | def forward(self, x):
class Bottleneck (line 70) | class Bottleneck(nn.Module):
method __init__ (line 73) | def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=N...
method forward (line 89) | def forward(self, x):
class ResNet (line 112) | class ResNet(nn.Module):
method __init__ (line 114) | def __init__(self, block, layers, num_classes=1000):
method _make_layer (line 141) | def _make_layer(self, block, planes, blocks, stride=1):
method forward (line 161) | def forward(self, x):
function resnet50_ibn_a (line 179) | def resnet50_ibn_a(pretrained=False, **kwargs):
function resnet101_ibn_a (line 192) | def resnet101_ibn_a(pretrained=False, **kwargs):
function remove_module_key (line 205) | def remove_module_key(state_dict):
FILE: hhcl/models/triplet.py
function euclidean_dist (line 8) | def euclidean_dist(x, y):
function cosine_dist (line 17) | def cosine_dist(x, y):
function _batch_hard (line 25) | def _batch_hard(mat_distance, mat_similarity, indice=False):
class TripletLoss (line 36) | class TripletLoss(nn.Module):
method __init__ (line 42) | def __init__(self, margin, normalize_feature=False):
method forward (line 48) | def forward(self, emb, label):
class SoftTripletLoss (line 65) | class SoftTripletLoss(nn.Module):
method __init__ (line 67) | def __init__(self, margin=None, normalize_feature=False):
method forward (line 72) | def forward(self, emb1, emb2, label):
FILE: hhcl/trainers.py
class Trainer (line 8) | class Trainer(object):
method __init__ (line 9) | def __init__(self, encoder, memory=None):
method train (line 14) | def train(self, epoch, data_loader, optimizer, print_freq=10, train_it...
method _parse_data (line 56) | def _parse_data(self, inputs):
method _forward (line 60) | def _forward(self, inputs):
FILE: hhcl/utils/__init__.py
function to_numpy (line 6) | def to_numpy(tensor):
function to_torch (line 15) | def to_torch(ndarray):
FILE: hhcl/utils/data/__init__.py
class IterLoader (line 7) | class IterLoader:
method __init__ (line 8) | def __init__(self, loader, length=None):
method __len__ (line 13) | def __len__(self):
method new_epoch (line 19) | def new_epoch(self):
method next (line 22) | def next(self):
FILE: hhcl/utils/data/base_dataset.py
class BaseDataset (line 5) | class BaseDataset(object):
method get_imagedata_info (line 10) | def get_imagedata_info(self, data):
method print_dataset_statistics (line 22) | def print_dataset_statistics(self):
method images_dir (line 26) | def images_dir(self):
class BaseImageDataset (line 30) | class BaseImageDataset(BaseDataset):
method print_dataset_statistics (line 35) | def print_dataset_statistics(self, train, query, gallery):
FILE: hhcl/utils/data/preprocessor.py
class Preprocessor (line 11) | class Preprocessor(Dataset):
method __init__ (line 12) | def __init__(self, dataset, root=None, transform=None, mutual=False):
method __len__ (line 19) | def __len__(self):
method __getitem__ (line 22) | def __getitem__(self, indices):
method _get_single_item (line 28) | def _get_single_item(self, index):
method _get_mutual_item (line 41) | def _get_mutual_item(self, index):
FILE: hhcl/utils/data/sampler.py
function No_index (line 14) | def No_index(a, b):
class RandomIdentitySampler (line 19) | class RandomIdentitySampler(Sampler):
method __init__ (line 20) | def __init__(self, data_source, num_instances):
method __len__ (line 29) | def __len__(self):
method __iter__ (line 32) | def __iter__(self):
class RandomMultipleGallerySampler (line 46) | class RandomMultipleGallerySampler(Sampler):
method __init__ (line 47) | def __init__(self, data_source, num_instances=4):
method __len__ (line 65) | def __len__(self):
method __iter__ (line 68) | def __iter__(self):
FILE: hhcl/utils/data/transforms.py
class RectScale (line 9) | class RectScale(object):
method __init__ (line 10) | def __init__(self, height, width, interpolation=Image.BILINEAR):
method __call__ (line 15) | def __call__(self, img):
class RandomSizedRectCrop (line 22) | class RandomSizedRectCrop(object):
method __init__ (line 23) | def __init__(self, height, width, interpolation=Image.BILINEAR):
method __call__ (line 28) | def __call__(self, img):
class RandomErasing (line 52) | class RandomErasing(object):
method __init__ (line 64) | def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=(0.4...
method __call__ (line 71) | def __call__(self, img):
FILE: hhcl/utils/faiss_rerank.py
function k_reciprocal_neigh (line 23) | def k_reciprocal_neigh(initial_rank, i, k1):
function compute_jaccard_distance (line 30) | def compute_jaccard_distance(target_features, k1=20, k2=6, print_flag=Tr...
FILE: hhcl/utils/faiss_utils.py
function swig_ptr_from_FloatTensor (line 6) | def swig_ptr_from_FloatTensor(x):
function swig_ptr_from_LongTensor (line 12) | def swig_ptr_from_LongTensor(x):
function search_index_pytorch (line 19) | def search_index_pytorch(index, x, k, D=None, I=None):
function search_raw_array_pytorch (line 44) | def search_raw_array_pytorch(res, xb, xq, k, D=None, I=None,
function index_init_gpu (line 92) | def index_init_gpu(ngpus, feat_dim):
function index_init_cpu (line 108) | def index_init_cpu(feat_dim):
FILE: hhcl/utils/logging.py
class Logger (line 8) | class Logger(object):
method __init__ (line 9) | def __init__(self, fpath=None):
method __del__ (line 16) | def __del__(self):
method __enter__ (line 19) | def __enter__(self):
method __exit__ (line 22) | def __exit__(self, *args):
method write (line 25) | def write(self, msg):
method flush (line 30) | def flush(self):
method close (line 36) | def close(self):
FILE: hhcl/utils/meters.py
class AverageMeter (line 4) | class AverageMeter(object):
method __init__ (line 7) | def __init__(self):
method reset (line 13) | def reset(self):
method update (line 19) | def update(self, val, n=1):
FILE: hhcl/utils/osutils.py
function mkdir_if_missing (line 6) | def mkdir_if_missing(dir_path):
FILE: hhcl/utils/rerank.py
function re_ranking (line 31) | def re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0...
FILE: hhcl/utils/serialization.py
function read_json (line 12) | def read_json(fpath):
function write_json (line 18) | def write_json(obj, fpath):
function save_checkpoint (line 24) | def save_checkpoint(state, is_best, fpath='checkpoint.pth.tar'):
function load_checkpoint (line 31) | def load_checkpoint(fpath):
function copy_state_dict (line 41) | def copy_state_dict(state_dict, model, strip=None):
Condensed preview — 41 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (138K chars).
[
{
"path": "README.md",
"chars": 4389,
"preview": "# HHCL-ReID \n[![Tweet](https://img.shield"
},
{
"path": "examples/test.py",
"chars": 4543,
"preview": "from __future__ import print_function, absolute_import\nimport argparse\nimport os.path as osp\nimport random\nimport numpy "
},
{
"path": "examples/train.py",
"chars": 13034,
"preview": "# -*- coding: utf-8 -*-\nfrom __future__ import print_function, absolute_import\nimport argparse\nimport os.path as osp\nimp"
},
{
"path": "hhcl/__init__.py",
"chars": 208,
"preview": "from __future__ import absolute_import\n\nfrom . import datasets\nfrom . import evaluation_metrics\nfrom . import models\nfro"
},
{
"path": "hhcl/datasets/__init__.py",
"chars": 1293,
"preview": "from __future__ import absolute_import\nimport warnings\n\nfrom .market1501 import Market1501\nfrom .msmt17 import MSMT17\nfr"
},
{
"path": "hhcl/datasets/celebreid.py",
"chars": 2988,
"preview": "from __future__ import print_function, absolute_import\nimport os.path as osp\nimport glob\nimport re\nfrom ..utils.data imp"
},
{
"path": "hhcl/datasets/dukemtmcreid.py",
"chars": 3121,
"preview": "import glob\nimport os.path as osp\nimport re\nfrom ..utils.data import BaseImageDataset\n\n\ndef process_dir(dir_path, relabe"
},
{
"path": "hhcl/datasets/market1501.py",
"chars": 3255,
"preview": "from __future__ import print_function, absolute_import\nimport os.path as osp\nimport glob\nimport re\nfrom ..utils.data imp"
},
{
"path": "hhcl/datasets/msmt17.py",
"chars": 2779,
"preview": "from __future__ import print_function, absolute_import\nimport os.path as osp\n\nimport glob\nimport re\nfrom ..utils.data im"
},
{
"path": "hhcl/datasets/personx.py",
"chars": 3033,
"preview": "from __future__ import print_function, absolute_import\nimport os.path as osp\nimport glob\nimport re\n\nfrom ..utils.data im"
},
{
"path": "hhcl/evaluation_metrics/__init__.py",
"chars": 167,
"preview": "from __future__ import absolute_import\n\nfrom .classification import accuracy\nfrom .ranking import cmc, mean_ap\n\n__all__ "
},
{
"path": "hhcl/evaluation_metrics/classification.py",
"chars": 604,
"preview": "from __future__ import absolute_import\n\nimport torch\nfrom ..utils import to_torch\n\n\ndef accuracy(output, target, topk=(1"
},
{
"path": "hhcl/evaluation_metrics/ranking.py",
"chars": 4126,
"preview": "from __future__ import absolute_import\nfrom collections import defaultdict\n\nimport numpy as np\nfrom sklearn.metrics impo"
},
{
"path": "hhcl/evaluators.py",
"chars": 4623,
"preview": "from __future__ import print_function, absolute_import\nimport time\nimport collections\nfrom collections import OrderedDic"
},
{
"path": "hhcl/models/__init__.py",
"chars": 1840,
"preview": "from __future__ import absolute_import\n\nfrom .resnet import *\nfrom .resnet_ibn import *\n\n__factory = {\n 'resnet18': r"
},
{
"path": "hhcl/models/cm.py",
"chars": 8792,
"preview": "import collections\nimport numpy as np\nfrom abc import ABC\nimport torch\nimport torch.nn.functional as F\nfrom torch import"
},
{
"path": "hhcl/models/dsbn.py",
"chars": 2680,
"preview": "import torch\nimport torch.nn as nn\n\n# Domain-specific BatchNorm\n\nclass DSBN2d(nn.Module):\n def __init__(self, planes)"
},
{
"path": "hhcl/models/kmeans.py",
"chars": 771,
"preview": "# Written by Yixiao Ge\n\nimport warnings\n\nimport faiss\nimport torch\n\nfrom ..utils import to_numpy, to_torch\n\n__all__ = [\""
},
{
"path": "hhcl/models/losses.py",
"chars": 12580,
"preview": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn import Parameter\n\n\nclass CrossEntropyLa"
},
{
"path": "hhcl/models/pooling.py",
"chars": 7009,
"preview": "# Credit to https://github.com/JDAI-CV/fast-reid/blob/master/fastreid/layers/pooling.py\nfrom abc import ABC\n\nimport torc"
},
{
"path": "hhcl/models/resnet.py",
"chars": 4516,
"preview": "from __future__ import absolute_import\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.nn import in"
},
{
"path": "hhcl/models/resnet_ibn.py",
"chars": 4049,
"preview": "from __future__ import absolute_import\n\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.nn import i"
},
{
"path": "hhcl/models/resnet_ibn_a.py",
"chars": 6595,
"preview": "import torch\nimport torch.nn as nn\nimport math\nimport torch.utils.model_zoo as model_zoo\n\n\n__all__ = ['ResNet', 'resnet5"
},
{
"path": "hhcl/models/triplet.py",
"chars": 3578,
"preview": "from __future__ import absolute_import\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\n\ndef euclidea"
},
{
"path": "hhcl/trainers.py",
"chars": 1848,
"preview": "from __future__ import print_function, absolute_import\nimport time\nimport torch\nimport torch.nn.functional as F\nfrom .ut"
},
{
"path": "hhcl/utils/__init__.py",
"chars": 594,
"preview": "from __future__ import absolute_import\n\nimport torch\n\n\ndef to_numpy(tensor):\n if torch.is_tensor(tensor):\n ret"
},
{
"path": "hhcl/utils/data/__init__.py",
"chars": 633,
"preview": "from __future__ import absolute_import\n\nfrom .base_dataset import BaseDataset, BaseImageDataset\nfrom .preprocessor impor"
},
{
"path": "hhcl/utils/data/base_dataset.py",
"chars": 1620,
"preview": "# encoding: utf-8\nimport numpy as np\n\n\nclass BaseDataset(object):\n \"\"\"\n Base class of reid dataset\n \"\"\"\n\n de"
},
{
"path": "hhcl/utils/data/preprocessor.py",
"chars": 1486,
"preview": "from __future__ import absolute_import\nimport os\nimport os.path as osp\nfrom torch.utils.data import DataLoader, Dataset\n"
},
{
"path": "hhcl/utils/data/sampler.py",
"chars": 3481,
"preview": "from __future__ import absolute_import\nfrom collections import defaultdict\nimport math\n\nimport numpy as np\nimport copy\ni"
},
{
"path": "hhcl/utils/data/transforms.py",
"chars": 3359,
"preview": "from __future__ import absolute_import\n\nfrom torchvision.transforms import *\nfrom PIL import Image\nimport random\nimport "
},
{
"path": "hhcl/utils/faiss_rerank.py",
"chars": 4846,
"preview": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\"\"\"\nCVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-i"
},
{
"path": "hhcl/utils/faiss_utils.py",
"chars": 3183,
"preview": "import os\nimport numpy as np\nimport faiss\nimport torch\n\ndef swig_ptr_from_FloatTensor(x):\n assert x.is_contiguous()\n "
},
{
"path": "hhcl/utils/logging.py",
"chars": 876,
"preview": "from __future__ import absolute_import\nimport os\nimport sys\n\nfrom .osutils import mkdir_if_missing\n\n\nclass Logger(object"
},
{
"path": "hhcl/utils/meters.py",
"chars": 497,
"preview": "from __future__ import absolute_import\n\n\nclass AverageMeter(object):\n \"\"\"Computes and stores the average and current "
},
{
"path": "hhcl/utils/osutils.py",
"chars": 214,
"preview": "from __future__ import absolute_import\nimport os\nimport errno\n\n\ndef mkdir_if_missing(dir_path):\n try:\n os.make"
},
{
"path": "hhcl/utils/rerank.py",
"chars": 4421,
"preview": "#!/usr/bin/env python2/python3\n# -*- coding: utf-8 -*-\n\"\"\"\nSource: https://github.com/zhunzhong07/person-re-ranking\nCrea"
},
{
"path": "hhcl/utils/serialization.py",
"chars": 1758,
"preview": "from __future__ import print_function, absolute_import\nimport json\nimport os.path as osp\nimport shutil\n\nimport torch\nfro"
},
{
"path": "requirements.txt",
"chars": 90,
"preview": "numpy\nsklearn\nCython\nh5py\npyzmq\npillow-simd\nsix\nscipy\nmatplotlib\nfaiss-gpu==1.6.3\neasydict"
},
{
"path": "run.sh",
"chars": 1295,
"preview": "### resnet50 ###\n# market1501\nCUDA_VISIBLE_DEVICES=0,1,2,3 python examples/train.py -b 256 -a resnet50 -d market1501 --i"
},
{
"path": "setup.py",
"chars": 420,
"preview": "from setuptools import setup, find_packages\n\n\nsetup(name='hhcl',\n version='1.0.0',\n install_requires=[\n "
}
]
About this extraction
This page contains the full source code of the bupt-ai-cz/HHCL-ReID GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 41 files (128.1 KB), approximately 34.0k tokens, and a symbol index with 248 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.