├── data └── README.md ├── utils ├── __init__.py ├── __pycache__ │ ├── util.cpython-37.pyc │ ├── losses.cpython-37.pyc │ ├── rerank.cpython-37.pyc │ ├── resnet.cpython-37.pyc │ ├── __init__.cpython-37.pyc │ ├── dataset.cpython-37.pyc │ ├── lr_adjust.cpython-37.pyc │ ├── ranking.cpython-37.pyc │ ├── evaluators.cpython-37.pyc │ ├── st_distribution.cpython-37.pyc │ └── evaluate_joint_sim.cpython-37.pyc ├── lr_adjust.py ├── evaluators.py ├── util.py ├── evaluate_joint_sim.py ├── rerank.py ├── dataset.py ├── ranking.py ├── resnet.py ├── st_distribution.py └── losses.py ├── imgs ├── results.png └── framework.png ├── performance.log ├── list_market ├── make_list_new.py └── make_list_new_test.py ├── README.md ├── market2duke_evaluat_joint_sim.py ├── duke2market_evaluat_joint_sim.py ├── loss_for_MSMT17 ├── gen_st_model_duke.py └── losses.py ├── market2duke_train.py └── duke2market_train.py /data/README.md: -------------------------------------------------------------------------------- 1 | data dir -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- 1 | from utils import * -------------------------------------------------------------------------------- /imgs/results.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/imgs/results.png -------------------------------------------------------------------------------- /imgs/framework.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/imgs/framework.png -------------------------------------------------------------------------------- /utils/__pycache__/util.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/util.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/losses.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/losses.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/rerank.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/rerank.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/resnet.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/resnet.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/__init__.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/__init__.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/dataset.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/dataset.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/lr_adjust.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/lr_adjust.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/ranking.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/ranking.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/evaluators.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/evaluators.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/st_distribution.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/st_distribution.cpython-37.pyc -------------------------------------------------------------------------------- /utils/__pycache__/evaluate_joint_sim.cpython-37.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/rika1024/JVTC/HEAD/utils/__pycache__/evaluate_joint_sim.cpython-37.pyc -------------------------------------------------------------------------------- /performance.log: -------------------------------------------------------------------------------- 1 | duke2market 2 | market test feature (19281, 512) 3 | performance based on visual similarity 4 | mAP: 0.6529, r1:0.8446, r5:0.9350, r10:0.9575, r20:0.9748 5 | performance based on joint similarity 6 | mAP: 0.6770, r1:0.8691, r5:0.9558, r10:0.9709, r20:0.9825 7 | 8 | 9 | 10 | 11 | market2duke 12 | duke test feature (19889, 512) 13 | performance based on visual similarity 14 | mAP: 0.5687, r1:0.7527, r5:0.8523, r10:0.8838, r20:0.9048 15 | performance based on joint similarity 16 | mAP: 0.6678, r1:0.8178, r5:0.9013, r10:0.9206, r20:0.9313 -------------------------------------------------------------------------------- /utils/lr_adjust.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | def SetLr(lr, optimizer): 5 | for g in optimizer.param_groups: 6 | g['lr'] = lr * g.get('lr_mult', 1) 7 | 8 | def StepLrUpdater(epoch, base_lr=0.01, gamma=0.1, step=[8,11]): 9 | if isinstance(step, int): 10 | return base_lr * (gamma**(epoch // step)) 11 | 12 | exp = len(step) 13 | for i, s in enumerate(step): 14 | if epoch < s: 15 | exp = i 16 | break 17 | return base_lr * gamma**exp 18 | 19 | 20 | def warmup_lr(cur_iters, warmup_iters=500, warmup_type='linear', warmup_ratio=1/3): 21 | 22 | if warmup_type == 'constant': 23 | warmup_lr = warmup_ratio 24 | elif warmup_type == 'linear': 25 | warmup_lr = 1 - (1 - cur_iters / warmup_iters) * (1 - warmup_ratio) 26 | elif warmup_type == 'exp': 27 | warmup_lr = warmup_ratio**(1 - cur_iters / warmup_iters) 28 | return warmup_lr 29 | 30 | -------------------------------------------------------------------------------- /list_market/make_list_new.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | from collections import defaultdict 3 | 4 | with open('list_market_train.txt') as f: 5 | lines = f.readlines() 6 | img_list = [i.split()[0] for i in lines] 7 | 8 | 9 | 10 | #id_seqs = defaultdict(list) 11 | seq_frames = defaultdict(list) 12 | for item in img_list: 13 | info = item.split('_') 14 | label = info[0] 15 | seq_id = info[1] 16 | frame = info[2] 17 | #id_seqs[label].append(seq_id) 18 | seq_frames[seq_id].append(int(frame)) 19 | 20 | seq_frames_accum = defaultdict(list) 21 | for key in seq_frames.keys(): 22 | seq_idx = int(key[-1]) 23 | 24 | if seq_idx == 1: 25 | seq_frames_accum[key] = 0#max(seq_frames[key]) 26 | else: 27 | accum = 0 28 | for i in range(1, seq_idx): 29 | accum_key = key[0:3] + str(i) 30 | accum = accum + max(seq_frames[accum_key]) 31 | seq_frames_accum[key] = accum 32 | 33 | # keys = seq_frames_accum.keys() 34 | # keys = list(keys) 35 | # keys.sort() 36 | # for key in keys: 37 | # print(key, seq_frames_accum[key], max(seq_frames[key])) 38 | 39 | 40 | file = open('list_market_train_new.txt', 'w') 41 | for line in lines: 42 | img_name = line.split()[0] 43 | seq_id = img_name.split('_')[1] 44 | 45 | line = line.split() 46 | frame_id = int(line[3]) 47 | new_line = "%s %s %s %d\n"%(line[0], line[1], line[2], frame_id+seq_frames_accum[seq_id]) 48 | file.write(new_line) -------------------------------------------------------------------------------- /list_market/make_list_new_test.py: -------------------------------------------------------------------------------- 1 | import os, sys 2 | from collections import defaultdict 3 | 4 | with open('list_market_train.txt') as f: 5 | lines = f.readlines() 6 | img_list = [i.split()[0] for i in lines] 7 | 8 | #id_seqs = defaultdict(list) 9 | seq_frames = defaultdict(list) 10 | for item in img_list: 11 | info = item.split('_') 12 | label = info[0] 13 | seq_id = info[1] 14 | frame = info[2] 15 | #id_seqs[label].append(seq_id) 16 | seq_frames[seq_id].append(int(frame)) 17 | 18 | seq_frames_accum = defaultdict(list) 19 | for key in seq_frames.keys(): 20 | seq_idx = int(key[-1]) 21 | 22 | if seq_idx == 1: 23 | seq_frames_accum[key] = 0#max(seq_frames[key]) 24 | else: 25 | accum = 0 26 | for i in range(1, seq_idx): 27 | accum_key = key[0:3] + str(i) 28 | accum = accum + max(seq_frames[accum_key]) 29 | seq_frames_accum[key] = accum 30 | 31 | # keys = seq_frames_accum.keys() 32 | # keys = list(keys) 33 | # keys.sort() 34 | # for key in keys: 35 | # print(key, seq_frames_accum[key], max(seq_frames[key])) 36 | 37 | 38 | with open('list_market_test.txt') as f: 39 | lines = f.readlines() 40 | 41 | file = open('list_market_test_new.txt', 'w') 42 | for line in lines: 43 | img_name = line.split()[0] 44 | #img_name = img_name.split('/') 45 | seq_id = img_name.split('/')[1].split('_')[1] 46 | 47 | line = line.split() 48 | frame_id = int(line[3]) 49 | new_line = "%s %s %s %d\n"%(line[0], line[1], line[2], frame_id+seq_frames_accum[seq_id]) 50 | file.write(new_line) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Introduction 2 | This is the code of [Joint Visual and Temporal Consistency for Unsupervised Domain Adaptive Person Re-Identification ECCV 2020](https://arxiv.org/pdf/2007.10854.pdf). 3 | 4 | ![](imgs/framework.png) 5 | 6 | ### Preparation 7 | 8 | #### Requirements: Python=3.6 and Pytorch>=1.0.0 9 | 10 | Please refer to [ECN](https://github.com/zhunzhong07/ECN) to prepare dataset, the file structure is as follow: 11 | 12 | ``` 13 | JVTC/data 14 | │ 15 | └───Market-1501 OR DukeMTMC-reID 16 | │ 17 | └───bounding_box_train 18 | │ 19 | └───bounding_box_test 20 | │ 21 | └───bounding_box_train_camstyle_merge 22 | | 23 | └───query 24 | ``` 25 | "bounding_box_train_camstyle_merge" dir merges the "bounding_box_train" and "bounding_box_train_camstyle" for convenience. 26 | 27 | ### Training and test 28 | We utilize 2 GTX-2080TI GPU for model training. 29 | ```Shell 30 | # Duke to Market-1501 training&evalution 31 | python duke2market_train.py 32 | 33 | # Duke to Market-1501 evalution with joint similarity 34 | python duke2market_evaluate_joint_sim.py 35 | 36 | # Market-1501 to Duke training&evalution 37 | python market2duke_train.py 38 | 39 | # Market-1501 to Duke evalution with joint similarity 40 | python market2duke_evaluate_joint_sim.py 41 | 42 | ``` 43 | 44 | ### Results 45 | 46 | ![](imgs/results.png) 47 | 48 | 49 | ### References 50 | 51 | - [1] Invariance Matters: Exemplar Memory for Domain Adaptive Person Re-identification. CVPR 2019. 52 | 53 | - [2] Spatial-temporal person re-identification. AAAI 2019. 54 | 55 | 56 | ### Citation 57 | 58 | If you find this code useful in your research, please consider citing: 59 | 60 | 61 | ### Contact me 62 | 63 | If you have any questions about this code or paper, please contact me at. 64 | 65 | [Jianing Li](https://kanei1024.github.io/) 66 | -------------------------------------------------------------------------------- /utils/evaluators.py: -------------------------------------------------------------------------------- 1 | import torch, cv2, os 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import numpy as np 5 | from .ranking import cmc, mean_ap 6 | 7 | #os.environ["CUDA_VISIBLE_DEVICES"] = "2" 8 | def extract_features(model, data_loader, select_set='market'): 9 | if select_set == 'market': 10 | query_num = 3368 11 | elif select_set =='duke': 12 | query_num = 2228 13 | 14 | print(select_set, "feature extraction start") 15 | model.eval() 16 | features = [] 17 | labels = [] 18 | cams = [] 19 | for i, (images, label, cam) in enumerate(data_loader): 20 | with torch.no_grad(): 21 | out = model(Variable(images).cuda()) 22 | features.append(out[1]) 23 | labels.append(label) 24 | cams.append(cam) 25 | 26 | features = torch.cat(features).cpu().numpy() 27 | labels = torch.cat(labels).cpu().numpy() 28 | cams = torch.cat(cams).cpu().numpy() 29 | #print('features', features.shape, labels.shape, cams.shape) 30 | 31 | query_labels = labels[0:query_num] 32 | query_cams = cams[0:query_num] 33 | 34 | gallery_labels = labels[query_num:] 35 | gallery_cams = cams[query_num:] 36 | 37 | query_features = features[0:query_num, :] 38 | gallery_features = features[query_num:, :] 39 | print("extraction done, feature shape:", np.shape(features)) 40 | 41 | return (query_features, query_labels, query_cams), (gallery_features, gallery_labels, gallery_cams) 42 | 43 | 44 | def evaluate_all(model, data_loader, select_set='market'): 45 | query, gallery = extract_features(model, data_loader, select_set=select_set) 46 | 47 | query_features, query_labels, query_cams = query 48 | gallery_features, gallery_labels, gallery_cams = gallery 49 | 50 | dist = np.zeros((query_features.shape[0], gallery_features.shape[0]), dtype = np.float64) 51 | for i in range(query_features.shape[0]): 52 | dist[i, :] = np.sum((gallery_features-query_features[i,:])**2, axis=1) 53 | 54 | mAP = mean_ap(dist, query_labels, gallery_labels, query_cams, gallery_cams) 55 | cmc_scores = cmc(dist, query_labels, gallery_labels, query_cams, gallery_cams, 56 | separate_camera_set=False, single_gallery_shot=False, first_match_break=True) 57 | print('mAP: %.4f, r1:%.4f, r5:%.4f, r10:%.4f, r20:%.4f'%(mAP, cmc_scores[0], cmc_scores[4], cmc_scores[9], cmc_scores[19])) 58 | 59 | return None -------------------------------------------------------------------------------- /utils/util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from sklearn.cluster import DBSCAN 3 | import torch 4 | from torch.nn import functional as F 5 | 6 | def l2_dist(fea_query, fea_gallery): 7 | dist = np.zeros((fea_query.shape[0], fea_gallery.shape[0]), dtype = np.float64) 8 | for i in range(fea_query.shape[0]): 9 | dist[i, :] = np.sum((fea_gallery-fea_query[i,:])**2, axis=1) 10 | return dist 11 | 12 | 13 | def cluster(dist, rho=1.6e-3): 14 | 15 | tri_mat = np.triu(dist, 1) 16 | tri_mat = tri_mat[np.nonzero(tri_mat)] 17 | tri_mat = np.sort(tri_mat,axis=None) 18 | top_num = np.round(rho*tri_mat.size).astype(int) 19 | eps = tri_mat[:top_num].mean() 20 | #low eps for training without source domain 21 | #eps = eps*0.8 22 | #print('eps in cluster: {:.3f}'.format(eps)) 23 | cluster = DBSCAN(eps=eps,min_samples=1,metric='precomputed', n_jobs=8) 24 | labels = cluster.fit_predict(dist) 25 | 26 | return labels 27 | 28 | 29 | def get_info(file_path): 30 | with open(file_path) as f: 31 | lines = f.readlines() 32 | #self.img_list = [os.path.join(dataset_dir, i.split()[0]) for i in lines] 33 | labels = [int(i.split()[1]) for i in lines] 34 | cam_ids = [int(i.split()[2]) for i in lines] 35 | frames = [int(i.split()[3]) for i in lines] 36 | 37 | return labels, cam_ids, frames 38 | 39 | def extract_fea_camtrans(model, loader): 40 | feas = [] 41 | for i, data in enumerate(loader, 1): 42 | #break 43 | with torch.no_grad(): 44 | image = data[0].cuda() 45 | 46 | batch_size = image.size(0) 47 | K = image.size(1) 48 | 49 | image = image.view(image.size(0)*image.size(1), image.size(2), image.size(3), image.size(4)) 50 | #image = Variable(image).cuda() 51 | out = model(image) 52 | fea = out[2] 53 | fea = fea.view(batch_size, K, -1) 54 | fea = fea.mean(dim=1) 55 | fea = F.normalize(fea) 56 | feas.append(fea) 57 | 58 | feas = torch.cat(feas) 59 | #print('duke_train_feas', feas.size()) 60 | return feas.cpu().numpy() 61 | 62 | def extract_fea_test(model, loader): 63 | feas = [] 64 | for i, data in enumerate(loader, 1): 65 | #break 66 | with torch.no_grad(): 67 | image = data[0].cuda() 68 | out = model(image) 69 | fea = out[1] 70 | feas.append(fea) 71 | 72 | feas = torch.cat(feas) 73 | #print('duke_train_feas', feas.size()) 74 | return feas.cpu().numpy() 75 | -------------------------------------------------------------------------------- /market2duke_evaluat_joint_sim.py: -------------------------------------------------------------------------------- 1 | import os, torch 2 | from torch.utils.data import DataLoader 3 | from torch.nn import functional as F 4 | import numpy as np 5 | from scipy.spatial.distance import cdist 6 | 7 | from utils.util import cluster, get_info 8 | from utils.util import extract_fea_camtrans, extract_fea_test 9 | from utils.resnet import resnet50 10 | from utils.dataset import imgdataset, imgdataset_camtrans 11 | from utils.rerank import re_ranking 12 | from utils.st_distribution import get_st_distribution 13 | from utils.evaluate_joint_sim import evaluate_joint 14 | 15 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 16 | 17 | dataset_path = 'data/dataset/' 18 | ann_file_train = 'list_duke/list_duke_train.txt' 19 | ann_file_test = 'list_duke/list_duke_test.txt' 20 | 21 | snapshot = 'snapshot/resnet50_market2duke_epoch00100.pth' 22 | 23 | num_cam = 8 24 | ########### DATASET ########### 25 | img_dir = dataset_path + 'DukeMTMC-reID/bounding_box_train_camstyle_merge/' 26 | train_dataset = imgdataset_camtrans(dataset_dir=img_dir, txt_path=ann_file_train, 27 | transformer='test', K=num_cam, num_cam=num_cam) 28 | train_loader = DataLoader(dataset=train_dataset, batch_size=1, shuffle=False, num_workers=4) 29 | 30 | img_dir = dataset_path + 'DukeMTMC-reID/' 31 | test_dataset = imgdataset(dataset_dir=img_dir, txt_path=ann_file_test, transformer='test') 32 | test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=4) 33 | 34 | ########### TEST ########### 35 | model, _ = resnet50(pretrained=snapshot, num_classes=751) 36 | model.cuda() 37 | model.eval() 38 | 39 | print('extract feature for training set') 40 | train_feas = extract_fea_camtrans(model, train_loader) 41 | _, cam_ids, frames = get_info(ann_file_train) 42 | 43 | print('generate spatial-temporal distribution') 44 | dist = cdist(train_feas, train_feas) 45 | dist = np.power(dist,2) 46 | #dist = re_ranking(original_dist=dist) 47 | labels = cluster(dist) 48 | num_ids = len(set(labels)) 49 | print('cluster id num:', num_ids) 50 | distribution = get_st_distribution(cam_ids, labels, frames, id_num=num_ids, cam_num=num_cam) 51 | 52 | print('extract feature for testing set') 53 | test_feas = extract_fea_test(model, test_loader) 54 | 55 | print('evaluation') 56 | evaluate_joint(test_fea=test_feas, st_distribute=distribution, ann_file=ann_file_test, select_set='duke') -------------------------------------------------------------------------------- /duke2market_evaluat_joint_sim.py: -------------------------------------------------------------------------------- 1 | import os, torch 2 | from torch.utils.data import DataLoader 3 | from torch.nn import functional as F 4 | import numpy as np 5 | from scipy.spatial.distance import cdist 6 | 7 | from utils.util import cluster, get_info 8 | from utils.util import extract_fea_camtrans, extract_fea_test 9 | from utils.resnet import resnet50 10 | from utils.dataset import imgdataset, imgdataset_camtrans 11 | from utils.rerank import re_ranking 12 | from utils.st_distribution import get_st_distribution 13 | from utils.evaluate_joint_sim import evaluate_joint 14 | 15 | os.environ["CUDA_VISIBLE_DEVICES"] = "1" 16 | 17 | dataset_path = 'data/dataset/' 18 | ann_file_train = 'list_market/list_market_train_new.txt' 19 | ann_file_test = 'list_market/list_market_test_new.txt' 20 | 21 | snapshot = 'snapshot/resnet50_duke2market_epoch00100.pth' 22 | 23 | num_cam = 6 24 | ########### DATASET ########### 25 | img_dir = dataset_path + 'Market-1501/bounding_box_train_camstyle_merge/' 26 | train_dataset = imgdataset_camtrans(dataset_dir=img_dir, txt_path=ann_file_train, 27 | transformer='test', K=num_cam, num_cam=num_cam) 28 | train_loader = DataLoader(dataset=train_dataset, batch_size=1, shuffle=False, num_workers=4) 29 | 30 | img_dir = dataset_path + 'Market-1501/' 31 | test_dataset = imgdataset(dataset_dir=img_dir, txt_path=ann_file_test, transformer='test') 32 | test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False, num_workers=4) 33 | 34 | ########### TEST ########### 35 | model, _ = resnet50(pretrained=snapshot, num_classes=702) 36 | model.cuda() 37 | model.eval() 38 | 39 | print('extract feature for training set') 40 | train_feas = extract_fea_camtrans(model, train_loader) 41 | _, cam_ids, frames = get_info(ann_file_train) 42 | 43 | print('generate spatial-temporal distribution') 44 | dist = cdist(train_feas, train_feas) 45 | dist = np.power(dist,2) 46 | #dist = re_ranking(original_dist=dist) 47 | labels = cluster(dist) 48 | num_ids = len(set(labels)) 49 | print('cluster id num:', num_ids) 50 | distribution = get_st_distribution(cam_ids, labels, frames, id_num=num_ids, cam_num=num_cam) 51 | 52 | print('extract feature for testing set') 53 | test_feas = extract_fea_test(model, test_loader) 54 | 55 | print('evaluation') 56 | evaluate_joint(test_fea=test_feas, st_distribute=distribution, ann_file=ann_file_test, select_set='market') 57 | -------------------------------------------------------------------------------- /utils/evaluate_joint_sim.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import numpy as np 5 | from torchvision import transforms 6 | 7 | from .rerank import re_ranking 8 | from .ranking import cmc, mean_ap 9 | from .st_distribution import joint_similarity 10 | from .util import get_info, l2_dist 11 | 12 | def compute_joint_dist(distribution, q_feas, g_feas, q_frames, g_frames, q_cams, g_cams): 13 | dists = [] 14 | for i in range(len(q_frames)): 15 | dist = joint_similarity( 16 | q_feas[i],q_cams[i],q_frames[i], 17 | g_feas, g_cams, g_frames, 18 | distribution) 19 | 20 | dist = np.expand_dims(dist, axis=0) 21 | dists.append(dist) 22 | #print(i, dist.shape) 23 | dists = np.concatenate(dists, axis=0) 24 | 25 | return dists 26 | 27 | def evaluate_joint(test_fea, st_distribute, ann_file, select_set='duke'): 28 | #fea_duke_test = np.load('duke_test_feas.npy') 29 | print(select_set, 'test feature', test_fea.shape) 30 | 31 | labels, cams, frames = get_info(ann_file) 32 | #print(len(cams)) 33 | #2228 for duke, 3368 for market, 11659 for msmt17 34 | if select_set == 'duke': 35 | query_num = 2228 36 | elif select_set == 'market': 37 | query_num = 3368 38 | 39 | query_labels = labels[0:query_num] 40 | query_cams = cams[0:query_num] 41 | query_frames = frames[0:query_num] 42 | query_features = test_fea[0:query_num, :] 43 | 44 | gallery_labels = labels[query_num:] 45 | gallery_cams = cams[query_num:] 46 | gallery_frames = frames[query_num:] 47 | gallery_features = test_fea[query_num:, :] 48 | 49 | 50 | dist = l2_dist(query_features, gallery_features) 51 | mAP = mean_ap(dist, query_labels, gallery_labels, query_cams, gallery_cams) 52 | cmc_scores = cmc(dist, query_labels, gallery_labels, query_cams, gallery_cams, 53 | separate_camera_set=False, single_gallery_shot=False, first_match_break=True) 54 | print('performance based on visual similarity') 55 | print('mAP: %.4f, r1:%.4f, r5:%.4f, r10:%.4f, r20:%.4f'%(mAP, cmc_scores[0], cmc_scores[4], cmc_scores[9], cmc_scores[19])) 56 | 57 | #st_distribute = np.load('distribution_duke_train.npy') 58 | dist = compute_joint_dist(st_distribute, 59 | query_features, gallery_features, 60 | query_frames, gallery_frames, 61 | query_cams, gallery_cams) 62 | 63 | mAP = mean_ap(dist, query_labels, gallery_labels, query_cams, gallery_cams) 64 | cmc_scores = cmc(dist, query_labels, gallery_labels, query_cams, gallery_cams, 65 | separate_camera_set=False, single_gallery_shot=False, first_match_break=True) 66 | print('performance based on joint similarity') 67 | print('mAP: %.4f, r1:%.4f, r5:%.4f, r10:%.4f, r20:%.4f'%(mAP, cmc_scores[0], cmc_scores[4], cmc_scores[9], cmc_scores[19])) 68 | -------------------------------------------------------------------------------- /utils/rerank.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # -*- coding: utf-8 -*- 3 | import numpy as np 4 | from scipy.spatial.distance import cdist 5 | 6 | 7 | def re_ranking(original_dist, k1=20, k2=6, lambda_value=0.3): 8 | 9 | all_num = original_dist.shape[0] 10 | 11 | 12 | euclidean_dist = original_dist 13 | gallery_num = original_dist.shape[0] #gallery_num=all_num 14 | 15 | #original_dist = original_dist - np.min(original_dist) 16 | original_dist = original_dist - np.min(original_dist,axis = 0) 17 | original_dist = np.transpose(original_dist/np.max(original_dist,axis = 0)) 18 | V = np.zeros_like(original_dist).astype(np.float16) 19 | initial_rank = np.argsort(original_dist).astype(np.int32) ## default axis=-1. 20 | 21 | print('Starting re_ranking...') 22 | for i in range(all_num): 23 | # k-reciprocal neighbors 24 | forward_k_neigh_index = initial_rank[i,:k1+1] ## k1+1 because self always ranks first. forward_k_neigh_index.shape=[k1+1]. forward_k_neigh_index[0] == i. 25 | backward_k_neigh_index = initial_rank[forward_k_neigh_index,:k1+1] ##backward.shape = [k1+1, k1+1]. For each ele in forward_k_neigh_index, find its rank k1 neighbors 26 | fi = np.where(backward_k_neigh_index==i)[0] 27 | k_reciprocal_index = forward_k_neigh_index[fi] ## get R(p,k) in the paper 28 | k_reciprocal_expansion_index = k_reciprocal_index 29 | for j in range(len(k_reciprocal_index)): 30 | candidate = k_reciprocal_index[j] 31 | candidate_forward_k_neigh_index = initial_rank[candidate,:int(np.around(k1/2))+1] 32 | candidate_backward_k_neigh_index = initial_rank[candidate_forward_k_neigh_index,:int(np.around(k1/2))+1] 33 | fi_candidate = np.where(candidate_backward_k_neigh_index == candidate)[0] 34 | candidate_k_reciprocal_index = candidate_forward_k_neigh_index[fi_candidate] 35 | if len(np.intersect1d(candidate_k_reciprocal_index,k_reciprocal_index))> 2/3*len(candidate_k_reciprocal_index): 36 | k_reciprocal_expansion_index = np.append(k_reciprocal_expansion_index,candidate_k_reciprocal_index) 37 | 38 | k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index) ## element-wise unique 39 | weight = np.exp(-original_dist[i,k_reciprocal_expansion_index]) 40 | V[i,k_reciprocal_expansion_index] = weight/np.sum(weight) 41 | #original_dist = original_dist[:query_num,] 42 | if k2 != 1: 43 | V_qe = np.zeros_like(V,dtype=np.float16) 44 | for i in range(all_num): 45 | V_qe[i,:] = np.mean(V[initial_rank[i,:k2],:],axis=0) 46 | V = V_qe 47 | del V_qe 48 | del initial_rank 49 | invIndex = [] 50 | for i in range(gallery_num): 51 | invIndex.append(np.where(V[:,i] != 0)[0]) #len(invIndex)=all_num 52 | 53 | jaccard_dist = np.zeros_like(original_dist,dtype = np.float16) 54 | 55 | 56 | for i in range(all_num): 57 | temp_min = np.zeros(shape=[1,gallery_num],dtype=np.float16) 58 | indNonZero = np.where(V[i,:] != 0)[0] 59 | indImages = [] 60 | indImages = [invIndex[ind] for ind in indNonZero] 61 | for j in range(len(indNonZero)): 62 | temp_min[0,indImages[j]] = temp_min[0,indImages[j]]+ np.minimum(V[i,indNonZero[j]],V[indImages[j],indNonZero[j]]) 63 | jaccard_dist[i] = 1-temp_min/(2-temp_min) 64 | 65 | pos_bool = (jaccard_dist < 0) 66 | jaccard_dist[pos_bool] = 0.0 67 | 68 | #np.save('dist_jaccard_temoral.npy', jaccard_dist) 69 | 70 | #return jaccard_dist 71 | if lambda_value == 0: 72 | return jaccard_dist 73 | else: 74 | final_dist = jaccard_dist*(1-lambda_value) + original_dist*lambda_value 75 | return final_dist 76 | 77 | -------------------------------------------------------------------------------- /utils/dataset.py: -------------------------------------------------------------------------------- 1 | import os, torch, random, cv2, math, glob 2 | import numpy as np 3 | from torch.utils import data 4 | from torchvision import transforms as T 5 | from PIL import Image 6 | from torch.nn import functional as F 7 | 8 | class RandomErasing(object): 9 | def __init__(self, EPSILON=0.5, mean=[0.485, 0.456, 0.406]): 10 | self.EPSILON = EPSILON 11 | self.mean = mean 12 | 13 | def __call__(self, img): 14 | 15 | if random.uniform(0, 1) > self.EPSILON: 16 | return img 17 | 18 | for attempt in range(100): 19 | area = img.size()[1] * img.size()[2] 20 | 21 | target_area = random.uniform(0.02, 0.2) * area 22 | aspect_ratio = random.uniform(0.3, 3) 23 | 24 | h = int(round(math.sqrt(target_area * aspect_ratio))) 25 | w = int(round(math.sqrt(target_area / aspect_ratio))) 26 | 27 | if w <= img.size()[2] and h <= img.size()[1]: 28 | x1 = random.randint(0, img.size()[1] - h) 29 | y1 = random.randint(0, img.size()[2] - w) 30 | img[0, x1:x1 + h, y1:y1 + w] = self.mean[0] 31 | img[1, x1:x1 + h, y1:y1 + w] = self.mean[1] 32 | img[2, x1:x1 + h, y1:y1 + w] = self.mean[2] 33 | return img 34 | 35 | return img 36 | 37 | normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 38 | train_transform = T.Compose([ 39 | T.Resize((256,128)), 40 | T.RandomHorizontalFlip(), 41 | T.ToTensor(), 42 | normalizer, 43 | RandomErasing(EPSILON=0.5) 44 | ]) 45 | 46 | test_transform = T.Compose([ 47 | T.Resize((256,128)), 48 | T.ToTensor(), 49 | normalizer ]) 50 | 51 | class imgdataset(data.Dataset): 52 | def __init__(self, dataset_dir, txt_path, transformer = 'train'): 53 | 54 | self.transform = train_transform if transformer == 'train' else test_transform 55 | with open(txt_path) as f: 56 | line = f.readlines() 57 | self.img_list = [os.path.join(dataset_dir, i.split()[0]) for i in line] 58 | self.label_list = [int(i.split()[1]) for i in line] 59 | self.cam_list = [int(i.split()[2]) for i in line] 60 | #self.cam_list = [int(i.split('c')[1][0]) for i in line] 61 | 62 | def __getitem__(self, index): 63 | im_path = self.img_list[index] 64 | image = Image.open(im_path) 65 | image = self.transform(image) 66 | return image, self.label_list[index], self.cam_list[index] 67 | 68 | def __len__(self): 69 | return len(self.label_list) 70 | 71 | 72 | class imgdataset_camtrans(data.Dataset): 73 | def __init__(self, dataset_dir, txt_path, transformer = 'train', num_cam=8, K=4): 74 | self.num_cam = num_cam 75 | self.transform = train_transform if transformer == 'train' else test_transform 76 | self.K = K 77 | with open(txt_path) as f: 78 | line = f.readlines() 79 | self.img_list = [os.path.join(dataset_dir, i.split()[0]) for i in line] 80 | self.label_list = [int(i.split()[1]) for i in line] 81 | #self.cam_list = [int(i.split('c')[1][0]) for i in line] 82 | self.cam_list = [int(i.split()[2]) for i in line] 83 | 84 | def __getitem__(self, index): 85 | im_path = self.img_list[index] 86 | camid = self.cam_list[index] 87 | cams = torch.randperm(self.num_cam) + 1 88 | 89 | imgs = [] 90 | for sel_cam in cams[0:self.K]: 91 | 92 | if sel_cam != camid: 93 | im_path_cam = im_path[:-4] + '_fake_' + str(camid) + 'to' + str(sel_cam.numpy()) + '.jpg' 94 | else: 95 | im_path_cam = im_path 96 | 97 | #print('im_path', camid, sel_cam,im_path_cam) 98 | image = Image.open(im_path_cam) 99 | image = self.transform(image) 100 | imgs.append(image.numpy()) 101 | #imgs.append(image) 102 | 103 | imgs = np.array(imgs, np.float32) 104 | imgs = torch.from_numpy(imgs).float() 105 | 106 | return imgs, self.label_list[index], index 107 | 108 | def __len__(self): 109 | return len(self.label_list) -------------------------------------------------------------------------------- /loss_for_MSMT17/gen_st_model_duke.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os, math 3 | 4 | def gaussian_func2(x, u, o=50): 5 | temp1 = 1.0 / (o * math.sqrt(2 * math.pi)) 6 | temp2 = -(np.power(x - u, 2)) / (2 * np.power(o, 2)) 7 | return temp1 * np.exp(temp2) 8 | 9 | def gauss_smooth2(arr,o): 10 | hist_num = len(arr) 11 | vect= np.zeros((hist_num,1)) 12 | for i in range(hist_num): 13 | vect[i,0]=i 14 | 15 | approximate_delta = 3*o # when x-u>approximate_delta, e.g., 6*o, the gaussian value is approximately equal to 0. 16 | gaussian_vect= gaussian_func2(vect,0,o) 17 | matrix = np.zeros((hist_num,hist_num)) 18 | for i in range(hist_num): 19 | k=0 20 | for j in range(i,hist_num): 21 | if k>approximate_delta: 22 | continue 23 | matrix[i][j]=gaussian_vect[j-i] 24 | k=k+1 25 | matrix = matrix+matrix.transpose() 26 | for i in range(hist_num): 27 | matrix[i][i]=matrix[i][i]/2 28 | 29 | xxx = np.dot(matrix,arr) 30 | return xxx 31 | 32 | 33 | #id_num = 1928 34 | def spatial_temporal_distribution(camera_id, labels, frames, id_num, cam_num=15): 35 | spatial_temporal_sum = np.zeros((id_num,cam_num)) 36 | spatial_temporal_count = np.zeros((id_num,cam_num)) 37 | eps = 0.0000001 38 | interval = 1.0 39 | 40 | for i in range(len(camera_id)): 41 | label_k = int(labels[i]) #### not in order, done 42 | cam_k = int(camera_id[i]-1) ##### ### ### ### ### ### ### ### ### ### ### ### # from 1, not 0 43 | frame_k = frames[i] 44 | 45 | spatial_temporal_sum[label_k][cam_k]=spatial_temporal_sum[label_k][cam_k]+frame_k 46 | spatial_temporal_count[label_k][cam_k] = spatial_temporal_count[label_k][cam_k] + 1 47 | spatial_temporal_avg = spatial_temporal_sum/(spatial_temporal_count+eps) # spatial_temporal_avg: 702 ids, 8cameras, center point 48 | 49 | distribution = np.zeros((cam_num,cam_num,3000)) 50 | for i in range(id_num): 51 | for j in range(cam_num-1): 52 | for k in range(j+1,cam_num): 53 | if spatial_temporal_count[i][j]==0 or spatial_temporal_count[i][k]==0: 54 | continue 55 | st_ij = spatial_temporal_avg[i][j] 56 | st_ik = spatial_temporal_avg[i][k] 57 | if st_ij>st_ik: 58 | diff = st_ij-st_ik 59 | hist_ = int(diff/interval) 60 | distribution[j][k][hist_] = distribution[j][k][hist_]+1 # [big][small] 61 | else: 62 | diff = st_ik-st_ij 63 | hist_ = int(diff/interval) 64 | distribution[k][j][hist_] = distribution[k][j][hist_]+1 65 | 66 | for i in range(id_num): 67 | for j in range(cam_num): 68 | if spatial_temporal_count[i][j] >1: 69 | 70 | frames_same_cam = [] 71 | for k in range(len(camera_id)): 72 | if labels[k]==i and camera_id[k]-1 ==j: 73 | frames_same_cam.append(frames[k]) 74 | frame_id_min = min(frames_same_cam) 75 | 76 | #print 'id, cam, len',i, j, len(frames_same_cam) 77 | for item in frames_same_cam: 78 | #if item != frame_id_min: 79 | diff = item - frame_id_min 80 | hist_ = int(diff/interval) 81 | #print item, frame_id_min, diff, hist_ 82 | distribution[j][j][hist_] = distribution[j][j][hist_] + spatial_temporal_count[i][j] 83 | 84 | smooth = 50 85 | for i in range(cam_num): 86 | for j in range(cam_num): 87 | #print("gauss "+str(i)+"->"+str(j)) 88 | distribution[i][j][:]=gauss_smooth2(distribution[i][j][:],smooth) 89 | 90 | sum_ = np.sum(distribution,axis=2) 91 | for i in range(cam_num): 92 | for j in range(cam_num): 93 | distribution[i][j][:]=distribution[i][j][:]/(sum_[i][j]+eps) 94 | 95 | return distribution # [to][from], to xxx camera, from xxx camera -------------------------------------------------------------------------------- /market2duke_train.py: -------------------------------------------------------------------------------- 1 | import os, torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | from torch.utils.data import DataLoader 5 | 6 | from utils.resnet import resnet50 7 | from utils.dataset import imgdataset, imgdataset_camtrans 8 | from utils.losses import Losses#, LocalLoss, GlobalLoss 9 | from utils.evaluators import evaluate_all 10 | from utils.lr_adjust import StepLrUpdater, SetLr 11 | 12 | os.environ["CUDA_VISIBLE_DEVICES"] = "2,3" 13 | ########### HYPER ########### 14 | base_lr = 0.01 15 | num_epoches = 100 16 | batch_size = 128 17 | K = 4 18 | num_cam = 8 19 | ########## DATASET ########### 20 | dataset_path = 'data/dataset/' 21 | src_dir = dataset_path + 'Market-1501/bounding_box_train/' 22 | tar_dir = dataset_path + 'DukeMTMC-reID/bounding_box_train_camstyle_merge/' 23 | tar_dir_test = dataset_path + 'DukeMTMC-reID/' 24 | 25 | src_annfile = 'list_market/list_market_train.txt' 26 | tar_annfile = 'list_duke/list_duke_train.txt' 27 | tar_annfile_test = 'list_duke/list_duke_test.txt' 28 | 29 | #resnet50: https://download.pytorch.org/models/resnet50-19c8e357.pth 30 | imageNet_pretrain = 'resnet50-19c8e357.pth' 31 | 32 | 33 | train_dataset = imgdataset(dataset_dir=src_dir, txt_path=src_annfile, transformer='train') 34 | train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True) 35 | 36 | train_dataset_t = imgdataset_camtrans(dataset_dir=tar_dir, txt_path=tar_annfile, 37 | transformer='train', num_cam=num_cam, K=K) 38 | train_loader_t = DataLoader(dataset=train_dataset_t, batch_size=int(batch_size/K), shuffle=True, num_workers=4, drop_last=True) 39 | 40 | test_dataset_t = imgdataset(dataset_dir=tar_dir_test, txt_path=tar_annfile_test, transformer='test') 41 | test_loader_t = DataLoader(dataset=test_dataset_t, batch_size=4, shuffle=False, num_workers=0) 42 | 43 | ########### MODEL ########### 44 | model, param = resnet50(pretrained=imageNet_pretrain, num_classes=751) 45 | model.cuda() 46 | model = nn.DataParallel(model)#, device_ids=[0,1]) 47 | 48 | losses = Losses(K=K, 49 | batch_size=batch_size, 50 | bank_size=len(train_dataset_t), 51 | ann_file=tar_annfile, 52 | cam_num=8) 53 | losses = losses.cuda() 54 | optimizer = torch.optim.SGD(param, lr=base_lr, momentum=0.9, weight_decay=5e-4, nesterov=True) 55 | 56 | ########### TRAIN ########### 57 | target_iter = iter(train_loader_t) 58 | for epoch in range(1, num_epoches+1): 59 | 60 | lr = StepLrUpdater(epoch, base_lr=base_lr, gamma=0.1, step=40) 61 | SetLr(lr, optimizer) 62 | 63 | print('-' * 10) 64 | print('Epoch [%d/%d], lr:%f'%(epoch, num_epoches, lr)) 65 | 66 | running_loss_src = 0.0 67 | running_loss_local = 0.0 68 | running_loss_global = 0.0 69 | 70 | if (epoch)%5 == 0: 71 | losses.reset_multi_label(epoch) 72 | 73 | model.train() 74 | for i, source_data in enumerate(train_loader, 1): 75 | try: 76 | target_data = next(target_iter) 77 | except: 78 | target_iter = iter(train_loader_t) 79 | target_data = next(target_iter) 80 | 81 | image_src = source_data[0].cuda() 82 | label_src = source_data[1].cuda() 83 | image_tar = target_data[0].cuda() 84 | image_tar = image_tar.view(-1, image_tar.size(2), image_tar.size(3), image_tar.size(4)) 85 | label_tar = target_data[2].cuda() 86 | 87 | x_src = model(image_src)[0] 88 | x_tar = model(image_tar)[2] 89 | 90 | loss_all= losses(x_src, label_src, x_tar, label_tar, epoch) 91 | loss, loss_s, loss_l, loss_g = loss_all 92 | 93 | 94 | running_loss_src += loss_s.mean().item() 95 | running_loss_local += loss_l.mean().item() 96 | running_loss_global += loss_g.mean().item() 97 | 98 | optimizer.zero_grad() 99 | loss.backward() 100 | optimizer.step() 101 | 102 | losses.update_memory(x_tar, label_tar, epoch=epoch) 103 | 104 | if i % 50 == 0: 105 | print(' iter: %3d/%d, loss src: %.3f, loss local: %.3f, loss global: %.3f'%(i, len(train_loader), running_loss_src/i, running_loss_local/i, running_loss_global/i)) 106 | print('Finish {} epoch\n'.format(epoch)) 107 | 108 | 109 | print('evaluation..') 110 | model.eval() 111 | evaluate_all(model, test_loader_t, select_set='duke') 112 | 113 | if hasattr(model, 'module'): 114 | model = model.module 115 | torch.save(model.state_dict(), 'snapshot/resnet50_market2duke_epoch%05d.pth'%(epoch)) 116 | print('save snapshot:','snapshot/resnet50_market2duke_epoch%05d.pth'%(epoch)) -------------------------------------------------------------------------------- /duke2market_train.py: -------------------------------------------------------------------------------- 1 | import os, torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | from torch.utils.data import DataLoader 5 | 6 | from utils.resnet import resnet50 7 | from utils.dataset import imgdataset, imgdataset_camtrans 8 | from utils.losses import Losses#, LocalLoss, GlobalLoss 9 | from utils.evaluators import evaluate_all 10 | from utils.lr_adjust import StepLrUpdater, SetLr 11 | 12 | os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" 13 | ########### HYPER ########### 14 | base_lr = 0.01 15 | num_epoches = 100 16 | batch_size = 128 17 | K = 4 18 | num_cam = 6 19 | ########## DATASET ########### 20 | dataset_path = 'data/dataset/' 21 | src_dir = dataset_path + 'DukeMTMC-reID/bounding_box_train/' 22 | tar_dir = dataset_path + 'Market-1501/bounding_box_train_camstyle_merge/' 23 | tar_dir_test = dataset_path + 'Market-1501/' 24 | 25 | src_annfile = 'list_duke/list_duke_train.txt' 26 | tar_annfile = 'list_market/list_market_train_new.txt' 27 | tar_annfile_test = 'list_market/list_market_test_new.txt' 28 | 29 | #resnet50: https://download.pytorch.org/models/resnet50-19c8e357.pth 30 | imageNet_pretrain = 'resnet50-19c8e357.pth' 31 | 32 | 33 | train_dataset = imgdataset(dataset_dir=src_dir, txt_path=src_annfile, transformer='train') 34 | train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True) 35 | 36 | train_dataset_t = imgdataset_camtrans(dataset_dir=tar_dir, txt_path=tar_annfile, 37 | transformer='train', num_cam=num_cam, K=K) 38 | train_loader_t = DataLoader(dataset=train_dataset_t, batch_size=int(batch_size/K), shuffle=True, num_workers=4, drop_last=True) 39 | 40 | test_dataset_t = imgdataset(dataset_dir=tar_dir_test, txt_path=tar_annfile_test, transformer='test') 41 | test_loader_t = DataLoader(dataset=test_dataset_t, batch_size=4, shuffle=False, num_workers=0) 42 | 43 | ########### MODEL ########### 44 | model, param = resnet50(pretrained=imageNet_pretrain, num_classes=702) 45 | model.cuda() 46 | model = nn.DataParallel(model)#, device_ids=[0,1]) 47 | 48 | losses = Losses(K=K, 49 | batch_size=batch_size, 50 | bank_size=len(train_dataset_t), 51 | ann_file=tar_annfile, 52 | cam_num=num_cam) 53 | losses = losses.cuda() 54 | optimizer = torch.optim.SGD(param, lr=base_lr, momentum=0.9, weight_decay=5e-4, nesterov=True) 55 | 56 | ########### TRAIN ########### 57 | target_iter = iter(train_loader_t) 58 | for epoch in range(1, num_epoches+1): 59 | 60 | lr = StepLrUpdater(epoch, base_lr=base_lr, gamma=0.1, step=40) 61 | SetLr(lr, optimizer) 62 | 63 | print('-' * 10) 64 | print('Epoch [%d/%d], lr:%f'%(epoch, num_epoches, lr)) 65 | 66 | running_loss_src = 0.0 67 | running_loss_local = 0.0 68 | running_loss_global = 0.0 69 | 70 | if (epoch)%5 == 0: 71 | losses.reset_multi_label(epoch) 72 | 73 | model.train() 74 | for i, source_data in enumerate(train_loader, 1): 75 | try: 76 | target_data = next(target_iter) 77 | except: 78 | target_iter = iter(train_loader_t) 79 | target_data = next(target_iter) 80 | 81 | image_src = source_data[0].cuda() 82 | label_src = source_data[1].cuda() 83 | image_tar = target_data[0].cuda() 84 | image_tar = image_tar.view(-1, image_tar.size(2), image_tar.size(3), image_tar.size(4)) 85 | label_tar = target_data[2].cuda() 86 | 87 | x_src = model(image_src)[0] 88 | x_tar = model(image_tar)[2] 89 | 90 | loss_all= losses(x_src, label_src, x_tar, label_tar, epoch) 91 | loss, loss_s, loss_l, loss_g = loss_all 92 | 93 | 94 | running_loss_src += loss_s.mean().item() 95 | running_loss_local += loss_l.mean().item() 96 | running_loss_global += loss_g.mean().item() 97 | 98 | optimizer.zero_grad() 99 | loss.backward() 100 | optimizer.step() 101 | 102 | losses.update_memory(x_tar, label_tar, epoch=epoch) 103 | 104 | if i % 50 == 0: 105 | print(' iter: %3d/%d, loss src: %.3f, loss local: %.3f, loss global: %.3f'%(i, len(train_loader), running_loss_src/i, running_loss_local/i, running_loss_global/i)) 106 | print('Finish {} epoch\n'.format(epoch)) 107 | 108 | 109 | print('evaluation..') 110 | model.eval() 111 | evaluate_all(model, test_loader_t, select_set='market') 112 | 113 | if hasattr(model, 'module'): 114 | model = model.module 115 | torch.save(model.state_dict(), 'snapshot/resnet50_duke2market_epoch%05d.pth'%(epoch)) 116 | print('save snapshot:','snapshot/resnet50_duke2market_epoch%05d.pth'%(epoch)) 117 | -------------------------------------------------------------------------------- /utils/ranking.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | from collections import defaultdict 3 | 4 | import numpy as np 5 | from sklearn.metrics import average_precision_score 6 | 7 | #from ..utils import to_numpy 8 | 9 | 10 | def _unique_sample(ids_dict, num): 11 | mask = np.zeros(num, dtype=np.bool) 12 | for _, indices in ids_dict.items(): 13 | i = np.random.choice(indices) 14 | mask[i] = True 15 | return mask 16 | 17 | 18 | def cmc(distmat, query_ids=None, gallery_ids=None, 19 | query_cams=None, gallery_cams=None, topk=100, 20 | separate_camera_set=False, 21 | single_gallery_shot=False, 22 | first_match_break=False): 23 | #distmat = to_numpy(distmat) 24 | m, n = distmat.shape 25 | # Fill up default values 26 | if query_ids is None: 27 | query_ids = np.arange(m) 28 | if gallery_ids is None: 29 | gallery_ids = np.arange(n) 30 | if query_cams is None: 31 | query_cams = np.zeros(m).astype(np.int32) 32 | if gallery_cams is None: 33 | gallery_cams = np.ones(n).astype(np.int32) 34 | # Ensure numpy array 35 | query_ids = np.asarray(query_ids) 36 | gallery_ids = np.asarray(gallery_ids) 37 | query_cams = np.asarray(query_cams) 38 | gallery_cams = np.asarray(gallery_cams) 39 | # Sort and find correct matches 40 | indices = np.argsort(distmat, axis=1) 41 | matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) 42 | # Compute CMC for each query 43 | ret = np.zeros(topk) 44 | num_valid_queries = 0 45 | for i in range(m): 46 | # Filter out the same id and same camera 47 | valid = ((gallery_ids[indices[i]] != query_ids[i]) | 48 | (gallery_cams[indices[i]] != query_cams[i])) 49 | if separate_camera_set: 50 | # Filter out samples from same camera 51 | valid &= (gallery_cams[indices[i]] != query_cams[i]) 52 | if not np.any(matches[i, valid]): continue 53 | if single_gallery_shot: 54 | repeat = 10 55 | gids = gallery_ids[indices[i][valid]] 56 | inds = np.where(valid)[0] 57 | ids_dict = defaultdict(list) 58 | for j, x in zip(inds, gids): 59 | ids_dict[x].append(j) 60 | else: 61 | repeat = 1 62 | for _ in range(repeat): 63 | if single_gallery_shot: 64 | # Randomly choose one instance for each id 65 | sampled = (valid & _unique_sample(ids_dict, len(valid))) 66 | index = np.nonzero(matches[i, sampled])[0] 67 | else: 68 | index = np.nonzero(matches[i, valid])[0] 69 | delta = 1. / (len(index) * repeat) 70 | for j, k in enumerate(index): 71 | if k - j >= topk: break 72 | if first_match_break: 73 | ret[k - j] += 1 74 | break 75 | ret[k - j] += delta 76 | num_valid_queries += 1 77 | if num_valid_queries == 0: 78 | raise RuntimeError("No valid query") 79 | return ret.cumsum() / num_valid_queries 80 | 81 | 82 | def mean_ap(distmat, query_ids=None, gallery_ids=None, 83 | query_cams=None, gallery_cams=None): 84 | #distmat = to_numpy(distmat) 85 | m, n = distmat.shape 86 | # Fill up default values 87 | if query_ids is None: 88 | query_ids = np.arange(m) 89 | if gallery_ids is None: 90 | gallery_ids = np.arange(n) 91 | if query_cams is None: 92 | query_cams = np.zeros(m).astype(np.int32) 93 | if gallery_cams is None: 94 | gallery_cams = np.ones(n).astype(np.int32) 95 | # Ensure numpy array 96 | query_ids = np.asarray(query_ids) 97 | gallery_ids = np.asarray(gallery_ids) 98 | query_cams = np.asarray(query_cams) 99 | gallery_cams = np.asarray(gallery_cams) 100 | # Sort and find correct matches 101 | indices = np.argsort(distmat, axis=1) 102 | matches = (gallery_ids[indices] == query_ids[:, np.newaxis]) 103 | # Compute AP for each query 104 | aps = [] 105 | for i in range(m): 106 | # Filter out the same id and same camera 107 | valid = ((gallery_ids[indices[i]] != query_ids[i]) | 108 | (gallery_cams[indices[i]] != query_cams[i])) #& gallery_ids[indices[i]] !=-1 109 | y_true = matches[i, valid] 110 | y_score = -distmat[i][indices[i]][valid] 111 | if not np.any(y_true): continue 112 | aps.append(average_precision_score(y_true, y_score)) 113 | if len(aps) == 0: 114 | raise RuntimeError("No valid query") 115 | return np.mean(aps) 116 | -------------------------------------------------------------------------------- /utils/resnet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math, torch 3 | import torch.utils.model_zoo as model_zoo 4 | from torch.nn import init 5 | from torch.nn import functional as F 6 | 7 | class Bottleneck(nn.Module): 8 | expansion = 4 9 | 10 | def __init__(self, inplanes, planes, stride=1, downsample=None): 11 | super(Bottleneck, self).__init__() 12 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 13 | self.bn1 = nn.BatchNorm2d(planes) 14 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 15 | padding=1, bias=False) 16 | self.bn2 = nn.BatchNorm2d(planes) 17 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 18 | self.bn3 = nn.BatchNorm2d(planes * 4) 19 | self.relu = nn.ReLU(inplace=True) 20 | self.downsample = downsample 21 | self.stride = stride 22 | 23 | def forward(self, x): 24 | residual = x 25 | 26 | out = self.conv1(x) 27 | out = self.bn1(out) 28 | out = self.relu(out) 29 | 30 | out = self.conv2(out) 31 | out = self.bn2(out) 32 | out = self.relu(out) 33 | 34 | out = self.conv3(out) 35 | out = self.bn3(out) 36 | 37 | if self.downsample is not None: 38 | residual = self.downsample(x) 39 | 40 | out += residual 41 | out = self.relu(out) 42 | 43 | return out 44 | 45 | 46 | class ResNet(nn.Module): 47 | 48 | def __init__(self, block, layers, num_classes=1000, train=True): 49 | self.inplanes = 64 50 | super(ResNet, self).__init__() 51 | self.istrain = train 52 | 53 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) 54 | self.bn1 = nn.BatchNorm2d(64) 55 | self.relu = nn.ReLU(inplace=True) 56 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 57 | self.layer1 = self._make_layer(block, 64, layers[0]) 58 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 59 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 60 | self.layer4 = self._make_layer(block, 512, layers[3], stride=1) 61 | #self.avgpool = nn.AvgPool2d((16,8), stride=1) 62 | 63 | self.num_features = 512 64 | self.feat = nn.Linear(512 * block.expansion, self.num_features) 65 | init.kaiming_normal_(self.feat.weight, mode='fan_out') 66 | init.constant_(self.feat.bias, 0) 67 | 68 | self.feat_bn = nn.BatchNorm1d(self.num_features) 69 | init.constant_(self.feat_bn.weight, 1) 70 | init.constant_(self.feat_bn.bias, 0) 71 | 72 | self.classifier = nn.Linear(self.num_features, num_classes) 73 | init.normal_(self.classifier.weight, std=0.001) 74 | init.constant_(self.classifier.bias, 0) 75 | 76 | 77 | for m in self.modules(): 78 | if isinstance(m, nn.Conv2d): 79 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 80 | m.weight.data.normal_(0, math.sqrt(2. / n)) 81 | elif isinstance(m, nn.BatchNorm2d): 82 | m.weight.data.fill_(1) 83 | m.bias.data.zero_() 84 | 85 | def _make_layer(self, block, planes, blocks, stride=1): 86 | downsample = None 87 | if stride != 1 or self.inplanes != planes * block.expansion: 88 | downsample = nn.Sequential( 89 | nn.Conv2d(self.inplanes, planes * block.expansion, 90 | kernel_size=1, stride=stride, bias=False), 91 | nn.BatchNorm2d(planes * block.expansion),) 92 | 93 | layers = [] 94 | layers.append(block(self.inplanes, planes, stride, downsample)) 95 | self.inplanes = planes * block.expansion 96 | for i in range(1, blocks): 97 | layers.append(block(self.inplanes, planes)) 98 | 99 | return nn.Sequential(*layers) 100 | 101 | def forward(self, x): 102 | x = self.conv1(x) 103 | x = self.bn1(x) 104 | x = self.relu(x) 105 | x = self.maxpool(x) 106 | 107 | x = self.layer1(x) 108 | x = self.layer2(x) 109 | x = self.layer3(x) 110 | x = self.layer4(x) 111 | 112 | x = F.avg_pool2d(x, x.size()[2:]) 113 | x = x.view(x.size(0), -1) 114 | 115 | x = self.feat(x) 116 | fea = self.feat_bn(x) 117 | fea_norm = F.normalize(fea) 118 | 119 | x = F.relu(fea) 120 | x = self.classifier(x) 121 | 122 | return x, fea_norm, fea 123 | 124 | def resnet50(pretrained=None, num_classes=1000, train=True): 125 | model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes, train) 126 | weight = torch.load(pretrained) 127 | static = model.state_dict() 128 | 129 | base_param = [] 130 | for name, param in weight.items(): 131 | if name not in static: 132 | continue 133 | if isinstance(param, nn.Parameter): 134 | param = param.data 135 | static[name].copy_(param) 136 | base_param.append(name) 137 | 138 | params = [] 139 | params_dict = dict(model.named_parameters()) 140 | for key, v in params_dict.items(): 141 | if key in base_param: 142 | params += [{ 'params':v, 'lr_mult':1}] 143 | else: 144 | #new parameter have larger learning rate 145 | params += [{ 'params':v, 'lr_mult':10}] 146 | 147 | return model, params -------------------------------------------------------------------------------- /utils/st_distribution.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import os, math 3 | 4 | 5 | def joint_similarity(qf,qc,qfr,gf,gc,gfr,distribution): 6 | query = qf 7 | score = np.dot(gf,query) 8 | gamma = 5 9 | 10 | interval = 100 11 | score_st = np.zeros(len(gc)) 12 | for i in range(len(gc)): 13 | if qfr>gfr[i]: 14 | diff = qfr-gfr[i] 15 | hist_ = int(diff/interval) 16 | pr = distribution[qc-1][gc[i]-1][hist_] 17 | else: 18 | diff = gfr[i]-qfr 19 | hist_ = int(diff/interval) 20 | pr = distribution[gc[i]-1][qc-1][hist_] 21 | score_st[i] = pr 22 | 23 | score = 1-1/(1+np.exp(-gamma*score))*1/(1+2*np.exp(-gamma*score_st)) 24 | return score 25 | 26 | def compute_joint_dist(distribution, q_feas, g_feas, q_frames, g_frames, q_cams, g_cams): 27 | dists = [] 28 | for i in range(len(q_frames)): 29 | dist = joint_similarity( 30 | q_feas[i],q_cams[i],q_frames[i], 31 | g_feas, g_cams, g_frames, 32 | distribution) 33 | 34 | dist = np.expand_dims(dist, axis=0) 35 | dists.append(dist) 36 | #print(i, dist.shape) 37 | dists = np.concatenate(dists, axis=0) 38 | 39 | return dists 40 | 41 | 42 | 43 | 44 | def gaussian_func(x, u, o=50): 45 | temp1 = 1.0 / (o * math.sqrt(2 * math.pi)) 46 | temp2 = -(np.power(x - u, 2)) / (2 * np.power(o, 2)) 47 | return temp1 * np.exp(temp2) 48 | 49 | def gauss_smooth(arr,o): 50 | hist_num = len(arr) 51 | vect= np.zeros((hist_num,1)) 52 | for i in range(hist_num): 53 | vect[i,0]=i 54 | 55 | approximate_delta = 3*o # when x-u>approximate_delta, e.g., 6*o, the gaussian value is approximately equal to 0. 56 | gaussian_vect= gaussian_func(vect,0,o) 57 | matrix = np.zeros((hist_num,hist_num)) 58 | for i in range(hist_num): 59 | k=0 60 | for j in range(i,hist_num): 61 | if k>approximate_delta: 62 | continue 63 | matrix[i][j]=gaussian_vect[j-i] 64 | k=k+1 65 | matrix = matrix+matrix.transpose() 66 | for i in range(hist_num): 67 | matrix[i][i]=matrix[i][i]/2 68 | 69 | xxx = np.dot(matrix,arr) 70 | return xxx 71 | 72 | def get_st_distribution(camera_id, labels, frames, id_num, cam_num=8): 73 | spatial_temporal_sum = np.zeros((id_num,cam_num)) 74 | spatial_temporal_count = np.zeros((id_num,cam_num)) 75 | eps = 0.0000001 76 | interval = 100.0 77 | 78 | for i in range(len(camera_id)): 79 | label_k = int(labels[i]) #### not in order, done 80 | cam_k = int(camera_id[i]-1) ##### ### ### ### ### ### ### ### ### ### ### ### # from 1, not 0 81 | frame_k = frames[i] 82 | 83 | spatial_temporal_sum[label_k][cam_k]=spatial_temporal_sum[label_k][cam_k]+frame_k 84 | spatial_temporal_count[label_k][cam_k] = spatial_temporal_count[label_k][cam_k] + 1 85 | spatial_temporal_avg = spatial_temporal_sum/(spatial_temporal_count+eps) # spatial_temporal_avg: 702 ids, 8cameras, center point 86 | 87 | distribution = np.zeros((cam_num,cam_num,3000)) 88 | for i in range(id_num): 89 | for j in range(cam_num-1): 90 | for k in range(j+1,cam_num): 91 | if spatial_temporal_count[i][j]==0 or spatial_temporal_count[i][k]==0: 92 | continue 93 | st_ij = spatial_temporal_avg[i][j] 94 | st_ik = spatial_temporal_avg[i][k] 95 | if st_ij>st_ik: 96 | diff = st_ij-st_ik 97 | hist_ = int(diff/interval) 98 | distribution[j][k][hist_] = distribution[j][k][hist_]+1 # [big][small] 99 | else: 100 | diff = st_ik-st_ij 101 | hist_ = int(diff/interval) 102 | distribution[k][j][hist_] = distribution[k][j][hist_]+1 103 | 104 | for i in range(id_num): 105 | for j in range(cam_num): 106 | if spatial_temporal_count[i][j] >1: 107 | 108 | frames_same_cam = [] 109 | for k in range(len(camera_id)): 110 | if labels[k]==i and camera_id[k]-1 ==j: 111 | frames_same_cam.append(frames[k]) 112 | frame_id_min = min(frames_same_cam) 113 | 114 | #print 'id, cam, len',i, j, len(frames_same_cam) 115 | for item in frames_same_cam: 116 | #if item != frame_id_min: 117 | diff = item - frame_id_min 118 | hist_ = int(diff/interval) 119 | #print item, frame_id_min, diff, hist_ 120 | distribution[j][j][hist_] = distribution[j][j][hist_] + spatial_temporal_count[i][j] 121 | 122 | smooth = 50 123 | for i in range(cam_num): 124 | for j in range(cam_num): 125 | #print("gauss "+str(i)+"->"+str(j)) 126 | distribution[i][j][:]=gauss_smooth(distribution[i][j][:],smooth) 127 | 128 | sum_ = np.sum(distribution,axis=2) 129 | for i in range(cam_num): 130 | for j in range(cam_num): 131 | distribution[i][j][:]=distribution[i][j][:]/(sum_[i][j]+eps) 132 | 133 | return distribution # [to][from], to xxx camera, from xxx camera 134 | 135 | 136 | 137 | -------------------------------------------------------------------------------- /utils/losses.py: -------------------------------------------------------------------------------- 1 | import torch, math, random 2 | import torch.nn.functional as F 3 | from torch import nn 4 | import numpy as np 5 | 6 | from scipy.spatial.distance import cdist 7 | 8 | from .rerank import re_ranking 9 | from .st_distribution import compute_joint_dist, get_st_distribution 10 | from .util import cluster 11 | 12 | 13 | class Losses(nn.Module): 14 | def __init__(self, K, batch_size, bank_size, ann_file, cam_num=8, beta1=0.1, beta2=0.05): 15 | super(Losses, self).__init__() 16 | 17 | self.loss_src = nn.CrossEntropyLoss()#.cuda() 18 | 19 | self.loss_local= LocalLoss(K=K, batch_size=batch_size, beta=beta1) 20 | self.loss_global = GlobalLoss(K=K, beta=beta2, bank_size=bank_size, cam_num=cam_num, ann_file=ann_file) 21 | 22 | def forward(self, x_src, label_src, x_tar, label_tar, epoch): 23 | 24 | loss_s = self.loss_src(x_src, label_src) 25 | 26 | loss_l = self.loss_local(x_tar) 27 | loss_g = self.loss_global(x_tar, label_tar) 28 | 29 | loss = loss_s + loss_l 30 | if epoch >= 10: 31 | loss = loss + loss_g * 0.2 32 | 33 | return loss, loss_s, loss_l, loss_g 34 | 35 | def reset_multi_label(self, epoch): 36 | print('Reset label on target dataset', epoch) 37 | if epoch >= 30: 38 | self.loss_global.reset_label_based_joint_smi() 39 | elif epoch >=10: 40 | self.loss_global.reset_label_based_visual_smi() 41 | 42 | def update_memory(self, x_tar, label_tar, epoch): 43 | self.loss_global.update(x_tar, label_tar, epoch=epoch) 44 | 45 | 46 | class LocalLoss(nn.Module): 47 | def __init__(self, K=4, batch_size=128, beta=0.1): 48 | super(LocalLoss, self).__init__() 49 | 50 | self.K = K 51 | self.beta = beta 52 | 53 | self.one_hot_label = [] 54 | for i in range(int(batch_size/K)): 55 | for j in range(K): 56 | self.one_hot_label.append(i) 57 | self.one_hot_label = torch.tensor(self.one_hot_label).cuda() 58 | 59 | def forward(self, x): 60 | 61 | v = x.view(int(x.size(0)/self.K), self.K, -1) 62 | v = v.mean(dim=1) 63 | v = F.normalize(v) 64 | 65 | x = F.normalize(x) 66 | x = x.mm(v.t())/self.beta 67 | loss = F.cross_entropy(x, self.one_hot_label) 68 | 69 | return loss 70 | 71 | # Global loss 72 | class GlobalLoss(nn.Module): 73 | def __init__(self, K=4, beta=0.05, cam_num=8, bank_size=114514, ann_file=None): 74 | super(GlobalLoss, self).__init__() 75 | 76 | self.K = K 77 | self.beta = beta 78 | self.cam_num = cam_num 79 | self.alpha = 0.01 80 | 81 | #self.bank = torch.zeros(bank_size, 512).cuda() 82 | self.bank = torch.rand(bank_size, 512).cuda() 83 | self.bank.requires_grad = False 84 | self.labels = torch.arange(0, bank_size).cuda() 85 | print('Memory bank size', self.bank.size()) 86 | 87 | with open(ann_file) as f: 88 | lines = f.readlines() 89 | #self.img_list = [os.path.join(dataset_dir, i.split()[0]) for i in lines] 90 | self.cam_ids = [int(i.split()[2]) for i in lines] 91 | self.frames = [int(i.split()[3]) for i in lines] 92 | 93 | print('dataset size:', len(self.cam_ids)) 94 | 95 | def reset_label_based_visual_smi(self): 96 | 97 | bank_fea = self.bank.cpu().data.numpy() 98 | #fea = fea.numpy() 99 | dist = cdist(bank_fea, bank_fea) 100 | dist = np.power(dist,2) 101 | print('Compute visual similarity') 102 | rerank_dist = re_ranking(original_dist=dist) 103 | 104 | labels = cluster(rerank_dist) 105 | num_ids = len(set(labels)) 106 | 107 | self.labels = torch.tensor(labels).cuda() 108 | print('Cluster class num based on visual similarity:', num_ids) 109 | 110 | def reset_label_based_joint_smi(self,): 111 | 112 | print('Compute distance based on visual similarity') 113 | 114 | bank_fea = self.bank.cpu().data.numpy() 115 | dist = cdist(bank_fea, bank_fea) 116 | dist = np.power(dist,2) 117 | 118 | #Jaccard distance for better cluster result 119 | dist = re_ranking(original_dist=dist) 120 | labels = cluster(dist) 121 | num_ids = len(set(labels)) 122 | 123 | print('update st distribution') 124 | st_distribute = get_st_distribution(self.cam_ids, labels, self.frames, id_num=num_ids, cam_num=self.cam_num) 125 | print('Compute distance based on joint similarity') 126 | st_dist = compute_joint_dist(st_distribute, 127 | bank_fea, bank_fea, 128 | self.frames, self.frames, 129 | self.cam_ids, self.cam_ids) 130 | 131 | #Jaccard distance for better cluster result 132 | st_dist = re_ranking(original_dist=st_dist, lambda_value=0.5) 133 | labels_st = cluster(st_dist) 134 | num_ids = len(set(labels_st)) 135 | 136 | print('Cluster class num based on joint similarity:', num_ids) 137 | self.labels = torch.tensor(labels_st).cuda() 138 | 139 | def forward(self, x, idx): 140 | 141 | w = x.view(int(x.size(0)/self.K), self.K, -1) 142 | w = w.mean(dim=1) 143 | w = F.normalize(w) 144 | x = F.normalize(x) 145 | 146 | label = w.mm(self.bank.t()) 147 | label = self.multi_class_label(idx) 148 | targets = [] 149 | for i in range(label.size(0)): 150 | for j in range(self.K): 151 | targets.append(label[i, :]) 152 | targets = torch.stack(targets).detach() 153 | 154 | x = x.mm(self.bank.t())/self.beta 155 | x = F.log_softmax(x, dim=1) 156 | loss = - (x * targets).sum(dim=1).mean(dim=0) 157 | 158 | #self.w = w.detach() 159 | #self.idx = idx.detach() 160 | return loss 161 | 162 | def update(self, x, idx, epoch=1): 163 | 164 | w = x.view(int(x.size(0)/self.K), self.K, -1) 165 | w = w.mean(dim=1) 166 | w = F.normalize(w).detach() 167 | 168 | momentum = min(self.alpha*epoch, 0.8) 169 | self.bank[idx] = w*(1-momentum) + self.bank[idx]*momentum 170 | # for i in range(self.w.size(0)): 171 | # self.bank[self.idx[i]] = self.w[i,:]*(1-momentum) + self.bank[self.idx[i]]*momentum 172 | # self.bank[self.idx[i]] = F.normalize(self.bank[self.idx[i]], dim=0) 173 | # #self.bank[idx[i]] /= self.bank[idx[i]].norm() 174 | 175 | def multi_class_label(self, index): 176 | batch_label = self.labels[index] 177 | target = (batch_label.unsqueeze(dim=1) == self.labels.t().unsqueeze(dim=0)).float() 178 | target = F.normalize(target, 1) 179 | #print target.size() 180 | return target -------------------------------------------------------------------------------- /loss_for_MSMT17/losses.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch import nn, autograd 4 | from torch.autograd import Variable, Function 5 | import numpy as np 6 | import math, rerank 7 | from sklearn.cluster import DBSCAN 8 | import gen_st_model_duke 9 | import random, math 10 | from scipy.spatial.distance import cdist 11 | 12 | class BatchLoss(nn.Module): 13 | def __init__(self, K=4, beta=0.05): 14 | super(BatchLoss, self).__init__() 15 | 16 | self.K = K 17 | self.beta = beta 18 | 19 | self.label = [] 20 | for i in range(32): 21 | for j in range(K): 22 | self.label.append(i) 23 | self.label = torch.tensor(self.label).cuda() 24 | 25 | def forward(self, x): 26 | 27 | w = x.view(int(x.size(0)/self.K), self.K, -1) 28 | w = w.mean(dim=1) 29 | w = F.normalize(w) 30 | 31 | x = F.normalize(x) 32 | x = x.mm(w.t())/self.beta 33 | loss = F.cross_entropy(x, self.label) 34 | 35 | return loss 36 | 37 | # Queue loss 38 | class QueueLoss(nn.Module): 39 | def __init__(self, K=4, beta=0.05, radio=0.2): 40 | super(QueueLoss, self).__init__() 41 | 42 | self.K = K 43 | self.beta = beta 44 | 45 | self.len = 32621 46 | #self.em = torch.zeros(self.len, 512).cuda() 47 | self.em = torch.rand(self.len, 512).cuda() 48 | 49 | 50 | self.em.requires_grad = False 51 | print self.em.size() 52 | self.alpha = 0.01 53 | 54 | self.labels = torch.arange(0,self.len).cuda() 55 | print'labels', self.labels.size() 56 | 57 | self.cam_ids = [] 58 | self.frames = [] 59 | self.time_slot = [] 60 | for item in open('list_MSMT/list_train.txt'): 61 | item = item.split()[0] 62 | camera = item.split('_')[2] 63 | frame = item.split('_')[4] 64 | self.cam_ids.append(int(camera)) 65 | self.frames.append(int(frame)) 66 | self.time_slot.append(item.split('_')[3]) 67 | #print frame[0:7] 68 | self.cam_ids = torch.tensor(self.cam_ids) 69 | self.frames = torch.tensor(self.frames) 70 | 71 | def reset_label(self, epoch): 72 | fea = self.em.cpu().data 73 | fea = fea.numpy() 74 | dist = cdist(fea, fea) 75 | dist = np.power(dist,2) 76 | print('Compute rerank dist, epoch', epoch) 77 | rerank_dist = rerank.re_ranking(original_dist=dist) 78 | 79 | rho = 1.6e-3 80 | tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2 81 | tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 82 | tri_mat = np.sort(tri_mat,axis=None) 83 | top_num = np.round(rho*tri_mat.size).astype(int) 84 | eps = tri_mat[:top_num].mean() 85 | eps = eps*0.8 86 | 87 | 88 | cluster = DBSCAN(eps=eps,min_samples=1,metric='precomputed', n_jobs=8) 89 | labels = cluster.fit_predict(rerank_dist) 90 | num_ids = len(set(labels)) - 1 91 | print 'eps in cluster: {:.3f}'.format(eps), 'num_ids', num_ids 92 | 93 | self.labels = torch.tensor(labels).cuda() 94 | 95 | 96 | 97 | 98 | def temporal_dist(self, qf,qc,qfr,gf,gc,gfr,distribution): 99 | query = qf 100 | score = np.dot(gf,query) 101 | alpha=5 102 | 103 | interval = 1.0 104 | score_st = np.zeros(len(gc)) 105 | for i in range(len(gc)): 106 | if qfr>gfr[i]: 107 | diff = qfr-gfr[i] 108 | hist_ = int(diff/interval) 109 | pr = distribution[qc-1][gc[i]-1][hist_] 110 | else: 111 | diff = gfr[i]-qfr 112 | hist_ = int(diff/interval) 113 | pr = distribution[gc[i]-1][qc-1][hist_] 114 | score_st[i] = pr 115 | score = 1-1/(1+np.exp(-alpha*score))*1/(1+2*np.exp(-alpha*score_st)) 116 | 117 | return score 118 | 119 | def compute_temporal_dist(self, labels, fea, idx): 120 | num_ids = len(set(labels)) 121 | frames = self.frames[idx] 122 | cam_ids = self.cam_ids[idx] 123 | 124 | distribution = gen_st_model_duke.spatial_temporal_distribution(cam_ids, labels, frames, id_num=num_ids, cam_num=15) 125 | 126 | dists = [] 127 | for i in range(len(frames)): 128 | dist = self.temporal_dist(fea[i], cam_ids[i], frames[i], fea, cam_ids, frames, distribution) 129 | dist = np.expand_dims(dist, axis=0) 130 | dists.append(dist) 131 | 132 | dists = np.concatenate(dists, axis=0) 133 | #print 'num_ids cluster', num_ids, len(idx), 'temporal dist shape', (dists.shape) 134 | return dists 135 | 136 | 137 | def reset_label_temporal(self, epoch): 138 | temporal_slot = ['0113morning','0113noon','0113afternoon', 139 | '0114morning','0114noon','0114afternoon', 140 | '0302morning','0302noon','0302afternoon', 141 | '0303morning','0303noon','0303afternoon'] 142 | 143 | 144 | label_sum = 0 145 | print('Compute rerank dist, epoch', epoch) 146 | for item in temporal_slot: 147 | time_divide = [] 148 | for i in range(len(self.time_slot)): 149 | if self.time_slot[i] == item: 150 | time_divide.append(i) 151 | 152 | divide_em = self.em[time_divide] 153 | fea = divide_em.data 154 | fea = fea.cpu().numpy() 155 | 156 | #dist = fea 157 | dist = cdist(fea, fea) 158 | dist = np.power(dist,2) 159 | rerank_dist = rerank.re_ranking(original_dist=dist) 160 | 161 | rho = 1.6e-3*3 162 | tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2 163 | tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 164 | tri_mat = np.sort(tri_mat,axis=None) 165 | top_num = np.round(rho*tri_mat.size).astype(int) 166 | eps = tri_mat[:top_num].mean()#*2 167 | #eps = eps*2 168 | cluster = DBSCAN(eps=eps,min_samples=1,metric='precomputed', n_jobs=8) 169 | labels = cluster.fit_predict(rerank_dist) 170 | num_ids = len(set(labels)) 171 | 172 | 173 | 174 | st_dist = self.compute_temporal_dist(labels, fea, time_divide) 175 | #print st_dist.shape 176 | rerank_st_dist = rerank.re_ranking(original_dist=st_dist)#, lambda_value=0.5) 177 | tri_mat = np.triu(rerank_st_dist, 1) # tri_mat.dim=2 178 | tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1 179 | tri_mat = np.sort(tri_mat,axis=None) 180 | top_num = np.round(rho*tri_mat.size).astype(int) 181 | eps = tri_mat[:top_num].mean() 182 | #print('eps in temporal cluster: {:.3f}'.format(eps)) 183 | cluster = DBSCAN(eps=eps,min_samples=1,metric='precomputed', n_jobs=8) 184 | labels = cluster.fit_predict(rerank_st_dist) 185 | num_ids = len(set(labels))# - 1 186 | #print('num_ids temporal cluster', num_ids) 187 | 188 | 189 | labels = labels + label_sum 190 | label_sum = label_sum + num_ids# + 1 191 | 192 | print 'eps: %f, num_ids: %d, label sum: %d'%(eps, num_ids, label_sum), item, rerank_st_dist.shape, st_dist.shape 193 | #print( label_sum, rerank_dist.shape) 194 | 195 | self.labels[time_divide] = torch.tensor(labels).cuda() 196 | 197 | def forward(self, x, idx): 198 | 199 | w = x.view(int(x.size(0)/self.K), self.K, -1) 200 | w = w.mean(dim=1) 201 | w = F.normalize(w) 202 | x = F.normalize(x) 203 | 204 | label = w.mm(self.em.t())#/self.beta 205 | label = self.smooth_rerank_label(idx) 206 | 207 | targets = [] 208 | for i in range(label.size(0)): 209 | for j in range(self.K): 210 | targets.append(label[i, :]) 211 | targets = torch.stack(targets).detach() 212 | 213 | x = x.mm(self.em.t())/self.beta 214 | x = F.log_softmax(x, dim=1) 215 | loss = - (x * targets).sum(dim=1).mean(dim=0) 216 | 217 | self.w = w.detach() 218 | self.idx = idx.detach() 219 | 220 | return loss 221 | 222 | def update(self, epoch=1): 223 | 224 | momentum = self.alpha*epoch 225 | if momentum > 0.8: 226 | momentum = 0.8 227 | for i in range(self.w.size(0)): 228 | self.em[self.idx[i]] = self.w[i,:]*(1-momentum) + self.em[self.idx[i]]*momentum 229 | self.em[self.idx[i]] = F.normalize(self.em[self.idx[i]],dim=0) 230 | #self.em[idx[i]] /= self.em[idx[i]].norm() 231 | 232 | 233 | def smooth_rerank_label(self, index): 234 | batch_label = self.labels[index] 235 | target = (batch_label.unsqueeze(dim=1) == self.labels.t().unsqueeze(dim=0)).float() 236 | target = F.normalize(target, 1) 237 | #print target.size() 238 | 239 | return target 240 | 241 | def smooth_hot(self, inputs, targets, k=6): 242 | # Sort 243 | _, index_sorted = torch.sort(inputs, dim=1, descending=True) 244 | 245 | ones_mat = torch.ones(targets.size(0), k).cuda()#to(self.device) 246 | targets = torch.unsqueeze(targets, 1) 247 | targets_onehot = torch.zeros(inputs.size()).cuda()#to(self.device) 248 | 249 | weights = F.softmax(ones_mat, dim=1) 250 | targets_onehot.scatter_(1, index_sorted[:, 0:k], ones_mat * weights) 251 | targets_onehot.scatter_(1, targets, float(1)) 252 | 253 | return targets_onehot 254 | 255 | 256 | # def smooth_hot(self, inputs, index, k=6): 257 | # _, index_sorted = torch.sort(inputs, dim=1, descending=True) 258 | # ones_mat = torch.ones(inputs.size(0), k).cuda() 259 | # targets_onehot = torch.zeros(inputs.size()).cuda() 260 | 261 | # inputs = (inputs > threshold).float() 262 | # targets_onehot.scatter_(1, index_sorted[:, 0:k], ones_mat) 263 | # targets_onehot = ((targets_onehot + inputs)>0).float() 264 | # #print targets_onehot.sum(dim=1), inputs.size() 265 | # targets_onehot = targets_onehot/targets_onehot.sum(dim=1, keepdim=True) 266 | 267 | # return targets_onehot --------------------------------------------------------------------------------