├── README.md ├── evaluate.py ├── evaluate_gpu.py ├── model.py ├── model.pyc ├── prepare.py ├── test.py └── train.py /README.md: -------------------------------------------------------------------------------- 1 | # ReID-PCB_RPP 2 | This repository tries to implement the paper:[Beyond Part Models: Person Retrieval with Refined Part Pooling](https://arxiv.org/abs/1711.09349) 3 | and I almost follow the training details at the origin paper. 4 | 5 | ## Market-1501 Results 6 | * The result of PCB is little higher than paper. 7 | * The result of PCB+RPP is little lower than paper. 8 | * I think if we do more augmentation with the training data, the performance will be better. 9 | 10 | | | Feature | mAP(%) | Rank-1(%) | 11 | | ------------- |:-------------:|:----------:| -----:| 12 | | PCB(paper) |G|77.4|92.3 13 | | PCB(paper) |H|77.3|92.4 14 | | PCB(ours) |G|**78.6**|**92.7** 15 | | PCB(ours) |H|78.5|92.1 16 | | PCB+RPP(paper) |G|81.6|93.8 17 | | PCB+RPP(paper) |H|81.0|93.1 18 | | PCB+RPP(ours) |G|**80.7**|**92.8** 19 | | PCB+RPP(ours) |H|79.8|92.4 20 | 21 | ## Prerequisites 22 | * Python 2.7 or 3.5 23 | * Pytorch 0.4 24 | * GPU Memory 25 | 26 | ## Getting started 27 | ### Dataset 28 | * Dwonload [Market1501 Dataset](http://www.liangzheng.org/Project/project_reid.html) and extract the files to the current folder 29 | * Change the *download_path* in the `prepare.py` and run `python prepare.py` to prepare the dataset 30 | 31 | ### Train 32 | Train the model: 33 | ``` 34 | python train.py --gpu_ids 0,1 --batchsize 64 --data_dir your_data_path --save_dir your_model_save_path --RPP True 35 | ``` 36 | * `--gpu_ids`: single GPU or mutil-GPUs 37 | * `--batchsize`: batch size 38 | * `--data_dir`: the path of the preparing data 39 | * `--save_dir`: the path to save the training model 40 | * `--RPP`: whether to train with RPP 41 | I train the model on the two GPUs with 64 batchsize 42 | ### Test 43 | Test the model(extract the features): 44 | ``` 45 | python test.py --gpu_ids 0 --which_epoch select the model --stage PCB or full --RPP True or False --feature_H True or False 46 | ``` 47 | * `--gpu_ids`: just single GPU 48 | * `--which_epoch`: select the i-th model 49 | * `--stage`: select the training model : PCB or full(PCB+RPP) 50 | * `--RPP`: if you choose the PCB model, RPP to False; if you choose the full model, RPP to True 51 | * `--feature_H`: whether to extract the low-dims features 52 | 53 | ### Evaluate 54 | Evaluate the model 55 | * Evaluate on GPU: 56 | ``` 57 | python evaluate_gpu.py --gpu_ids 0 --reslut_mat the path of the features mat 58 | ``` 59 | * Evaluate on CPU: 60 | ``` 61 | python evaluate_gpu.py --reslut_mat the path of the features mat 62 | ``` 63 | 64 | ## Reference resources 65 | Thanks to the [layumi/Person_reID_baseline_pytorch](https://github.com/layumi/Person_reID_baseline_pytorch/blob/b9465f16a3b94300ceee8884120226804beee224/README.md). 66 | This repository only implements the part of PCB, I make some modifications on it and then add the RPP part. 67 | 68 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import scipy.io 2 | import torch 3 | import numpy as np 4 | import time 5 | import os 6 | import argparse 7 | 8 | parser = argparse.ArgumentParser(description='Evaluating') 9 | parser.add_argument('--result_mat', default='./result/RPP_G_result.mat', type=str, help='save result dir') 10 | args = parser.parse_args() 11 | ####################################################################### 12 | # Evaluate 13 | def evaluate(qf,ql,qc,gf,gl,gc): 14 | query = qf 15 | score = np.dot(gf,query) 16 | # predict index 17 | index = np.argsort(score) #from small to large 18 | index = index[::-1] 19 | #index = index[0:2000] 20 | # good index 21 | query_index = np.argwhere(gl==ql) 22 | camera_index = np.argwhere(gc==qc) 23 | 24 | good_index = np.setdiff1d(query_index, camera_index, assume_unique=True) 25 | junk_index1 = np.argwhere(gl==-1) 26 | junk_index2 = np.intersect1d(query_index, camera_index) 27 | # junk_index2 = camera_index 28 | junk_index = np.append(junk_index2, junk_index1) #.flatten()) 29 | 30 | CMC_tmp = compute_mAP(index, good_index, junk_index) 31 | return CMC_tmp 32 | 33 | 34 | def compute_mAP(index, good_index, junk_index): 35 | ap = 0 36 | cmc = torch.IntTensor(len(index)).zero_() 37 | if good_index.size==0: # if empty 38 | cmc[0] = -1 39 | return ap,cmc 40 | 41 | # remove junk_index 42 | mask = np.in1d(index, junk_index, invert=True) 43 | index = index[mask] 44 | 45 | # find good_index index 46 | ngood = len(good_index) 47 | mask = np.in1d(index, good_index) 48 | rows_good = np.argwhere(mask==True) 49 | rows_good = rows_good.flatten() 50 | 51 | cmc[rows_good[0]:] = 1 52 | for i in range(ngood): 53 | d_recall = 1.0/ngood 54 | precision = (i+1)*1.0/(rows_good[i]+1) 55 | if rows_good[i]!=0: 56 | old_precision = i*1.0/rows_good[i] 57 | else: 58 | old_precision=1.0 59 | ap = ap + d_recall*(old_precision + precision)/2 60 | 61 | return ap, cmc 62 | 63 | ###################################################################### 64 | result = scipy.io.loadmat(args.result_mat) 65 | query_feature = result['query_f'] 66 | query_cam = result['query_cam'][0] 67 | query_label = result['query_label'][0] 68 | gallery_feature = result['gallery_f'] 69 | gallery_cam = result['gallery_cam'][0] 70 | gallery_label = result['gallery_label'][0] 71 | 72 | multi = os.path.isfile('multi_query.mat') 73 | 74 | 75 | CMC = torch.IntTensor(len(gallery_label)).zero_() 76 | ap = 0.0 77 | #print(query_label) 78 | for i in range(len(query_label)): 79 | ap_tmp, CMC_tmp = evaluate(query_feature[i],query_label[i],query_cam[i],gallery_feature,gallery_label,gallery_cam) 80 | if CMC_tmp[0]==-1: 81 | continue 82 | CMC = CMC + CMC_tmp 83 | ap += ap_tmp 84 | print(i, CMC_tmp[0]) 85 | 86 | CMC = CMC.float() 87 | CMC = CMC/len(query_label) #average CMC 88 | print('Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f'%(CMC[0],CMC[4],CMC[9],ap/len(query_label))) 89 | 90 | -------------------------------------------------------------------------------- /evaluate_gpu.py: -------------------------------------------------------------------------------- 1 | import scipy.io 2 | import torch 3 | import numpy as np 4 | import time 5 | import os 6 | import argparse 7 | 8 | ###################################################################### 9 | # 10 | parser = argparse.ArgumentParser(description='Evaluating') 11 | parser.add_argument('--result_mat', default='./result/PCB_H_result.mat', type=str, help='save result dir') 12 | parser.add_argument('--gpu_ids', default='0', type=int,help='gpu_ids: e.g. 0') 13 | args = parser.parse_args() 14 | 15 | # setGPU 16 | torch.cuda.set_device(args.gpu_ids) 17 | ####################################################################### 18 | # Evaluate 19 | def evaluate(qf,ql,qc,gf,gl,gc): 20 | query = qf.view(-1,1) 21 | # print(query.shape) 22 | score = torch.mm(gf,query) 23 | score = score.squeeze(1).cpu() 24 | score = score.numpy() 25 | # predict index 26 | index = np.argsort(score) #from small to large 27 | index = index[::-1] 28 | # index = index[0:2000] 29 | # good index 30 | query_index = np.argwhere(gl==ql) 31 | camera_index = np.argwhere(gc==qc) 32 | 33 | good_index = np.setdiff1d(query_index, camera_index, assume_unique=True) 34 | junk_index1 = np.argwhere(gl==-1) 35 | junk_index2 = np.intersect1d(query_index, camera_index) 36 | junk_index = np.append(junk_index2, junk_index1) #.flatten()) 37 | 38 | CMC_tmp = compute_mAP(index, good_index, junk_index) 39 | return CMC_tmp 40 | 41 | 42 | def compute_mAP(index, good_index, junk_index): 43 | ap = 0 44 | cmc = torch.IntTensor(len(index)).zero_() 45 | if good_index.size==0: # if empty 46 | cmc[0] = -1 47 | return ap,cmc 48 | 49 | # remove junk_indexevaluate_gpu.py:33 50 | mask = np.in1d(index, junk_index, invert=True) 51 | index = index[mask] 52 | 53 | # find good_index index 54 | ngood = len(good_index) 55 | mask = np.in1d(index, good_index) 56 | rows_good = np.argwhere(mask==True) 57 | rows_good = rows_good.flatten() 58 | 59 | cmc[rows_good[0]:] = 1 60 | for i in range(ngood): 61 | d_recall = 1.0/ngood 62 | precision = (i+1)*1.0/(rows_good[i]+1) 63 | if rows_good[i]!=0: 64 | old_precision = i*1.0/rows_good[i] 65 | else: 66 | old_precision=1.0 67 | ap = ap + d_recall*(old_precision + precision)/2 68 | 69 | return ap, cmc 70 | 71 | ###################################################################### 72 | result = scipy.io.loadmat(args.result_mat) 73 | query_feature = torch.FloatTensor(result['query_f']) 74 | query_cam = result['query_cam'][0] 75 | query_label = result['query_label'][0] 76 | gallery_feature = torch.FloatTensor(result['gallery_f']) 77 | gallery_cam = result['gallery_cam'][0] 78 | gallery_label = result['gallery_label'][0] 79 | 80 | 81 | 82 | query_feature = query_feature.cuda() 83 | gallery_feature = gallery_feature.cuda() 84 | 85 | print(query_feature.shape) 86 | CMC = torch.IntTensor(len(gallery_label)).zero_() 87 | ap = 0.0 88 | #print(query_label) 89 | for i in range(len(query_label)): 90 | ap_tmp, CMC_tmp = evaluate(query_feature[i],query_label[i],query_cam[i],gallery_feature,gallery_label,gallery_cam) 91 | if CMC_tmp[0]==-1: 92 | continue 93 | CMC = CMC + CMC_tmp 94 | ap += ap_tmp 95 | #print(i, CMC_tmp[0]) 96 | 97 | CMC = CMC.float() 98 | CMC = CMC/len(query_label) #average CMC 99 | print('Rank@1:%f Rank@5:%f Rank@10:%f mAP:%f'%(CMC[0],CMC[4],CMC[9],ap/len(query_label))) 100 | 101 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.nn import init 4 | from torchvision import models 5 | from torch.autograd import Variable 6 | 7 | 8 | ###################################################################### 9 | def weights_init_kaiming(m): 10 | classname = m.__class__.__name__ 11 | if classname.find('Conv2d') != -1: 12 | init.kaiming_normal_(m.weight.data, mode='fan_out', nonlinearity='relu') 13 | elif classname.find('Linear') != -1: 14 | init.kaiming_normal(m.weight.data, a=0, mode='fan_out') 15 | init.constant(m.bias.data, 0.0) 16 | elif classname.find('BatchNorm1d') != -1: 17 | init.normal(m.weight.data, 1.0, 0.02) 18 | init.constant(m.bias.data, 0.0) 19 | elif classname.find('BatchNorm2d') != -1: 20 | init.constant(m.weight.data, 1) 21 | init.constant(m.bias.data, 0) 22 | 23 | 24 | def weights_init_classifier(m): 25 | classname = m.__class__.__name__ 26 | if classname.find('Linear') != -1: 27 | init.normal(m.weight.data, std=0.001) 28 | init.constant(m.bias.data, 0.0) 29 | 30 | 31 | # Defines the new fc layer and classification layer 32 | # |--Linear--|--bn--|--relu--|--Linear--| 33 | class ClassBlock(nn.Module): 34 | def __init__(self, input_dim, class_num, relu=True, num_bottleneck=512): 35 | super(ClassBlock, self).__init__() 36 | add_block = [] 37 | 38 | add_block += [nn.Conv2d(input_dim, num_bottleneck, kernel_size=1, bias=False)] 39 | add_block += [nn.BatchNorm2d(num_bottleneck)] 40 | if relu: 41 | #add_block += [nn.LeakyReLU(0.1)] 42 | add_block += [nn.ReLU(inplace=True)] 43 | add_block = nn.Sequential(*add_block) 44 | add_block.apply(weights_init_kaiming) 45 | 46 | classifier = [] 47 | classifier += [nn.Linear(num_bottleneck, class_num)] 48 | classifier = nn.Sequential(*classifier) 49 | classifier.apply(weights_init_classifier) 50 | 51 | self.add_block = add_block 52 | self.classifier = classifier 53 | 54 | def forward(self, x): 55 | x = self.add_block(x) 56 | x = torch.squeeze(x) 57 | x = self.classifier(x) 58 | return x 59 | 60 | # Define the RPP layers 61 | class RPP(nn.Module): 62 | def __init__(self): 63 | super(RPP, self).__init__() 64 | self.part = 6 65 | add_block = [] 66 | add_block += [nn.Conv2d(2048, 6, kernel_size=1, bias=False)] 67 | add_block = nn.Sequential(*add_block) 68 | add_block.apply(weights_init_kaiming) 69 | 70 | norm_block = [] 71 | norm_block += [nn.BatchNorm2d(2048)] 72 | norm_block += [nn.ReLU(inplace=True)] 73 | # norm_block += [nn.LeakyReLU(0.1, inplace=True)] 74 | norm_block = nn.Sequential(*norm_block) 75 | norm_block.apply(weights_init_kaiming) 76 | 77 | self.add_block = add_block 78 | self.norm_block = norm_block 79 | self.softmax = nn.Softmax(dim=1) 80 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 81 | 82 | 83 | def forward(self, x): 84 | w = self.add_block(x) 85 | p = self.softmax(w) 86 | y = [] 87 | for i in range(self.part): 88 | p_i = p[:, i, :, :] 89 | p_i = torch.unsqueeze(p_i, 1) 90 | y_i = torch.mul(x, p_i) 91 | y_i = self.norm_block(y_i) 92 | y_i = self.avgpool(y_i) 93 | y.append(y_i) 94 | 95 | f = torch.cat(y, 2) 96 | return f 97 | 98 | 99 | # Part Model proposed in Yifan Sun etal. (2018) 100 | class PCB(nn.Module): 101 | def __init__(self, class_num): 102 | super(PCB, self).__init__() 103 | 104 | self.part = 6 105 | # resnet50 106 | resnet = models.resnet50(pretrained=True) 107 | # remove the final downsample 108 | resnet.layer4[0].downsample[0].stride = (1, 1) 109 | resnet.layer4[0].conv2.stride = (1, 1) 110 | modules = list(resnet.children())[:-2] 111 | self.backbone = nn.Sequential(*modules) 112 | self.avgpool = nn.AdaptiveAvgPool2d((self.part, 1)) 113 | self.dropout = nn.Dropout(p=0.5) 114 | 115 | # define 6 classifiers 116 | self.classifiers = nn.ModuleList() 117 | for i in range(self.part): 118 | self.classifiers.append(ClassBlock(2048, class_num, True, 256)) 119 | 120 | def forward(self, x): 121 | x = self.backbone(x) 122 | x = self.avgpool(x) 123 | x = self.dropout(x) 124 | part = {} 125 | predict = {} 126 | # get six part feature batchsize*2048*6 127 | for i in range(self.part): 128 | part[i] = x[:, :, i, :] 129 | part[i] = torch.unsqueeze(part[i], 3) 130 | # print part[i].shape 131 | predict[i] = self.classifiers[i](part[i]) 132 | 133 | y = [] 134 | for i in range(self.part): 135 | y.append(predict[i]) 136 | return y 137 | 138 | def convert_to_rpp(self): 139 | self.avgpool = RPP() 140 | return self 141 | 142 | 143 | class PCB_test(nn.Module): 144 | def __init__(self, model, featrue_H=False): 145 | super(PCB_test, self).__init__() 146 | self.part = 6 147 | self.featrue_H = featrue_H 148 | self.backbone = model.backbone 149 | self.avgpool = model.avgpool 150 | self.classifiers = nn.ModuleList() 151 | for i in range(self.part): 152 | self.classifiers.append(model.classifiers[i].add_block) 153 | 154 | def forward(self, x): 155 | x = self.backbone(x) 156 | x = self.avgpool(x) 157 | 158 | if self.featrue_H: 159 | part = {} 160 | predict = {} 161 | # get six part feature batchsize*2048*6 162 | for i in range(self.part): 163 | part[i] = x[:, :, i, :] 164 | part[i] = torch.unsqueeze(part[i], 3) 165 | predict[i] = self.classifiers[i](part[i]) 166 | 167 | y = [] 168 | for i in range(self.part): 169 | y.append(predict[i]) 170 | x = torch.cat(y, 2) 171 | f = x.view(x.size(0), x.size(1), x.size(2)) 172 | return f 173 | 174 | 175 | # debug model structure 176 | 177 | net = PCB(751) 178 | net = net.convert_to_rpp() 179 | print(net) 180 | input = Variable(torch.FloatTensor(8, 3, 7, 7)) 181 | output = net(input) 182 | # print(output[0].shape) 183 | -------------------------------------------------------------------------------- /model.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xiaoccer/ReID-PCB_RPP/00462d044864f23581820b9696ca02afe04391e7/model.pyc -------------------------------------------------------------------------------- /prepare.py: -------------------------------------------------------------------------------- 1 | import os 2 | from shutil import copyfile 3 | 4 | # You only need to change this line to your dataset download path 5 | download_path = './Market' 6 | 7 | if not os.path.isdir(download_path): 8 | print('please change the download_path') 9 | 10 | save_path = download_path + '/pytorch' 11 | if not os.path.isdir(save_path): 12 | os.mkdir(save_path) 13 | #----------------------------------------- 14 | #query 15 | query_path = download_path + '/query' 16 | query_save_path = download_path + '/pytorch/query' 17 | if not os.path.isdir(query_save_path): 18 | os.mkdir(query_save_path) 19 | 20 | for root, dirs, files in os.walk(query_path, topdown=True): 21 | for name in files: 22 | if not name[-3:]=='jpg': 23 | continue 24 | ID = name.split('_') 25 | src_path = query_path + '/' + name 26 | dst_path = query_save_path + '/' + ID[0] 27 | if not os.path.isdir(dst_path): 28 | os.mkdir(dst_path) 29 | copyfile(src_path, dst_path + '/' + name) 30 | 31 | #----------------------------------------- 32 | #multi-query 33 | query_path = download_path + '/gt_bbox' 34 | query_save_path = download_path + '/pytorch/multi-query' 35 | if not os.path.isdir(query_save_path): 36 | os.mkdir(query_save_path) 37 | 38 | for root, dirs, files in os.walk(query_path, topdown=True): 39 | for name in files: 40 | if not name[-3:]=='jpg': 41 | continue 42 | ID = name.split('_') 43 | src_path = query_path + '/' + name 44 | dst_path = query_save_path + '/' + ID[0] 45 | if not os.path.isdir(dst_path): 46 | os.mkdir(dst_path) 47 | copyfile(src_path, dst_path + '/' + name) 48 | 49 | #----------------------------------------- 50 | #gallery 51 | gallery_path = download_path + '/bounding_box_test' 52 | gallery_save_path = download_path + '/pytorch/gallery' 53 | if not os.path.isdir(gallery_save_path): 54 | os.mkdir(gallery_save_path) 55 | 56 | for root, dirs, files in os.walk(gallery_path, topdown=True): 57 | for name in files: 58 | if not name[-3:]=='jpg': 59 | continue 60 | ID = name.split('_') 61 | src_path = gallery_path + '/' + name 62 | dst_path = gallery_save_path + '/' + ID[0] 63 | if not os.path.isdir(dst_path): 64 | os.mkdir(dst_path) 65 | copyfile(src_path, dst_path + '/' + name) 66 | 67 | #--------------------------------------- 68 | #train_all 69 | train_path = download_path + '/bounding_box_train' 70 | train_save_path = download_path + '/pytorch/train_all' 71 | if not os.path.isdir(train_save_path): 72 | os.mkdir(train_save_path) 73 | 74 | for root, dirs, files in os.walk(train_path, topdown=True): 75 | for name in files: 76 | if not name[-3:]=='jpg': 77 | continue 78 | ID = name.split('_') 79 | src_path = train_path + '/' + name 80 | dst_path = train_save_path + '/' + ID[0] 81 | if not os.path.isdir(dst_path): 82 | os.mkdir(dst_path) 83 | copyfile(src_path, dst_path + '/' + name) 84 | 85 | 86 | #--------------------------------------- 87 | #train_val 88 | train_path = download_path + '/bounding_box_train' 89 | train_save_path = download_path + '/pytorch/train' 90 | val_save_path = download_path + '/pytorch/val' 91 | if not os.path.isdir(train_save_path): 92 | os.mkdir(train_save_path) 93 | os.mkdir(val_save_path) 94 | 95 | for root, dirs, files in os.walk(train_path, topdown=True): 96 | for name in files: 97 | if not name[-3:]=='jpg': 98 | continue 99 | ID = name.split('_') 100 | src_path = train_path + '/' + name 101 | dst_path = train_save_path + '/' + ID[0] 102 | if not os.path.isdir(dst_path): 103 | os.mkdir(dst_path) 104 | dst_path = val_save_path + '/' + ID[0] #first image is used as val image 105 | os.mkdir(dst_path) 106 | copyfile(src_path, dst_path + '/' + name) 107 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import print_function, division 4 | 5 | import argparse 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | from torch.optim import lr_scheduler 10 | from torch.autograd import Variable 11 | import numpy as np 12 | import torchvision 13 | from torchvision import datasets, models, transforms 14 | import time 15 | import os 16 | import scipy.io 17 | from model import PCB, PCB_test 18 | 19 | ###################################################################### 20 | # Options 21 | # -------- 22 | parser = argparse.ArgumentParser(description='Testing') 23 | parser.add_argument('--gpu_ids', default='0', type=str,help='gpu_ids: e.g. 0 0,1,2 0,2') 24 | parser.add_argument('--which_epoch', default='last', type=str, help='0,1,2,3...or last') 25 | parser.add_argument('--test_dir', default='/home/brl/USER/fzc/PCB_RPP/Market/pytorch',type=str, 26 | help='./test_data') 27 | parser.add_argument('--model_dir', default='./model/', type=str, help='save model dir') 28 | parser.add_argument('--result_dir', default='./result/', type=str, help='save result dir') 29 | parser.add_argument('--stage', default='PCB', type=str, help='save model path') 30 | parser.add_argument('--RPP', action='store_true', help='use RPP', default=False) 31 | parser.add_argument('--batchsize', default=32, type=int, help='batch_size') 32 | parser.add_argument('--feature_H', action='store_true', help='extract feature_H', default=True) 33 | args = parser.parse_args() 34 | 35 | # --- 36 | str_ids = args.gpu_ids.split(',') 37 | stage = args.stage 38 | test_dir = args.test_dir 39 | model_dir = args.model_dir 40 | result_dir = args.result_dir 41 | feature_H = args.feature_H 42 | 43 | # setGPU 44 | gpu_ids = [] 45 | for str_id in str_ids: 46 | id = int(str_id) 47 | if id >=0: 48 | gpu_ids.append(id) 49 | if len(gpu_ids)>0: 50 | torch.cuda.set_device(gpu_ids[0]) 51 | use_gpu = torch.cuda.is_available() 52 | 53 | ###################################################################### 54 | # Load Data 55 | # --------- 56 | # 57 | # We will use torchvision and torch.utils.data packages for loading the 58 | # data. 59 | # 60 | 61 | data_transforms = transforms.Compose([ 62 | transforms.Resize((384, 128), interpolation=3), 63 | transforms.ToTensor(), 64 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 65 | ]) 66 | 67 | 68 | data_dir = test_dir 69 | image_datasets = {x: datasets.ImageFolder( os.path.join(data_dir,x), data_transforms) for x in ['gallery','query']} 70 | dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batchsize, 71 | shuffle=False, num_workers=8) for x in ['gallery', 'query']} 72 | 73 | class_names = image_datasets['query'].classes 74 | 75 | ###################################################################### 76 | # Load model 77 | #--------------------------- 78 | def load_network(network): 79 | save_path = os.path.join(args.model_dir, stage, 'net_%s.pth'%args.which_epoch) 80 | network.load_state_dict(torch.load(save_path)) 81 | return network 82 | 83 | 84 | ###################################################################### 85 | # Extract feature 86 | # ---------------------- 87 | # 88 | # Extract feature from a trained model. 89 | # 90 | def fliplr(img): 91 | '''flip horizontal''' 92 | inv_idx = torch.arange(img.size(3)-1,-1,-1).long() # N x C x H x W 93 | img_flip = img.index_select(3, inv_idx) 94 | return img_flip 95 | 96 | def extract_feature(model,dataloaders): 97 | features = torch.FloatTensor() 98 | count = 0 99 | for data in dataloaders: 100 | img, label = data 101 | n, c, h, w = img.size() 102 | count += n 103 | print(count) 104 | if args.feature_H: 105 | ff = torch.FloatTensor(n, 256, 6).zero_() # we have six parts 106 | else: 107 | ff = torch.FloatTensor(n, 2048, 6).zero_() # we have six parts 108 | for i in range(2): 109 | if(i == 1): 110 | img = fliplr(img) 111 | input_img = Variable(img.cuda()) 112 | outputs = model(input_img) 113 | f = outputs.data.cpu() 114 | ff = ff+f 115 | 116 | # feature size (n,2048,6) 117 | # 1. To treat every part equally, I calculate the norm for every 2048-dim part feature. 118 | # 2. To keep the cosine score==1, sqrt(6) is added to norm the whole feature (2048*6). 119 | fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) * np.sqrt(6) 120 | # fnorm = torch.norm(ff, p=2, dim=1, keepdim=True) 121 | ff = ff.div(fnorm.expand_as(ff)) 122 | ff = ff.view(ff.size(0), -1) 123 | features = torch.cat((features, ff), 0) 124 | 125 | return features 126 | 127 | def get_id(img_path): 128 | camera_id = [] 129 | labels = [] 130 | for path, v in img_path: 131 | #filename = path.split('/')[-1] 132 | filename = os.path.basename(path) 133 | label = filename[0:4] 134 | camera = filename.split('c')[1] 135 | if label[0:2]=='-1': 136 | labels.append(-1) 137 | else: 138 | labels.append(int(label)) 139 | camera_id.append(int(camera[0])) 140 | return camera_id, labels 141 | 142 | gallery_path = image_datasets['gallery'].imgs 143 | query_path = image_datasets['query'].imgs 144 | 145 | gallery_cam,gallery_label = get_id(gallery_path) 146 | query_cam,query_label = get_id(query_path) 147 | 148 | ###################################################################### 149 | # Load Collected data Trained model 150 | print('-------test-----------') 151 | 152 | 153 | model_structure = PCB(751) 154 | if args.RPP: 155 | model_structure = model_structure.convert_to_rpp() 156 | 157 | model = load_network(model_structure) 158 | 159 | model = PCB_test(model, feature_H) 160 | 161 | # Change to test mode 162 | model = model.eval() 163 | if use_gpu: 164 | model = model.cuda() 165 | 166 | # Extract feature 167 | gallery_feature = extract_feature(model, dataloaders['gallery']) 168 | query_feature = extract_feature(model, dataloaders['query']) 169 | 170 | 171 | # Save to Matlab for check 172 | result = {'gallery_f':gallery_feature.numpy(),'gallery_label':gallery_label,'gallery_cam':gallery_cam,'query_f':query_feature.numpy(),'query_label':query_label,'query_cam':query_cam} 173 | if not os.path.isdir(args.result_dir): 174 | os.mkdir(args.result_dir) 175 | if args.RPP: 176 | save_pre = '/RPP_' 177 | else: 178 | save_pre = '/PCB_' 179 | if args.feature_H: 180 | save_pre += 'H_' 181 | else: 182 | save_pre += 'G_' 183 | scipy.io.savemat(args.result_dir +save_pre+ 'result.mat', result) 184 | 185 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import print_function, division 4 | 5 | import argparse 6 | import torch 7 | import torch.nn as nn 8 | import torch.optim as optim 9 | from torch.optim import lr_scheduler 10 | from torch.autograd import Variable 11 | import numpy as np 12 | import torchvision 13 | from torchvision import datasets, models, transforms 14 | import matplotlib 15 | 16 | matplotlib.use('agg') 17 | import matplotlib.pyplot as plt 18 | from PIL import Image 19 | import time 20 | import os 21 | from model import PCB 22 | import json 23 | from shutil import copyfile 24 | 25 | version = torch.__version__ 26 | 27 | ###################################################################### 28 | # Options # 29 | parser = argparse.ArgumentParser(description='Training') 30 | parser.add_argument('--gpu_ids', default='0, 1', type=str, help='gpu_ids: e.g. 0 0,1,2 0,2') 31 | parser.add_argument('--save_dir', default='./model/', type=str, help='save model dir') 32 | parser.add_argument('--data_dir', default='/home/brl/USER/fzc/PCB_RPP/Market/pytorch', type=str, help='training dir path') 33 | parser.add_argument('--train_all', action='store_true', default=True, help='use all training data') 34 | parser.add_argument('--batchsize', default=64, type=int, help='batch_size') 35 | parser.add_argument('--RPP', action='store_true', help='use RPP', default=True) 36 | args = parser.parse_args() 37 | 38 | # setGPU # 39 | use_gpu = torch.cuda.is_available() 40 | str_ids = args.gpu_ids.split(',') 41 | gpu_ids = [] 42 | for str_id in str_ids: 43 | gid = int(str_id) 44 | if gid >= 0: 45 | gpu_ids.append(gid) 46 | 47 | is_parallel_train = False 48 | if len(gpu_ids) > 1: 49 | is_parallel_train = True 50 | 51 | def get_net(is_parallel, net): 52 | return net.module if is_parallel else net 53 | 54 | # setSeed # 55 | seed = 1994 56 | torch.manual_seed(seed) 57 | torch.cuda.manual_seed(seed) 58 | 59 | # ------ # 60 | data_dir = args.data_dir 61 | save_dir = args.save_dir 62 | if not os.path.isdir(save_dir): 63 | os.mkdir(save_dir) 64 | 65 | ###################################################################### 66 | # Load Data 67 | # --------- 68 | # 69 | 70 | transform_train_list = [ 71 | transforms.Resize((384, 128), interpolation=3), 72 | transforms.RandomHorizontalFlip(), 73 | transforms.ToTensor(), 74 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 75 | ] 76 | transform_val_list = [ 77 | transforms.Resize(size=(384, 128), interpolation=3), # Image.BICUBIC 78 | transforms.ToTensor(), 79 | transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) 80 | ] 81 | 82 | data_transforms = { 83 | 'train': transforms.Compose(transform_train_list), 84 | 'val': transforms.Compose(transform_val_list), 85 | } 86 | 87 | train_all = '' 88 | if args.train_all: 89 | train_all = '_all' 90 | 91 | image_datasets = {} 92 | image_datasets['train'] = datasets.ImageFolder(os.path.join(data_dir, 'train' + train_all), 93 | data_transforms['train']) 94 | image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'), 95 | data_transforms['val']) 96 | 97 | dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=args.batchsize, 98 | shuffle=True, num_workers=8) # 8 workers may work faster 99 | for x in ['train', 'val']} 100 | dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} 101 | class_names = image_datasets['train'].classes 102 | 103 | ###################################################################### 104 | # Training the model 105 | # ------------------ 106 | # 107 | # Now, let's write a general function to train a model. Here, we will 108 | # illustrate: 109 | # 110 | # - Scheduling the learning rate 111 | # - Saving the best model 112 | # 113 | # In the following, parameter ``scheduler`` is an LR scheduler object from 114 | # ``torch.optim.lr_scheduler``. 115 | 116 | def train_model(model, criterion, optimizer, scheduler, log_file, stage, num_epochs=25): 117 | since = time.time() 118 | 119 | best_model_wts = model.state_dict() 120 | last_model_wts = model.state_dict() 121 | best_acc = 0.0 122 | 123 | for epoch in range(num_epochs): 124 | print('Epoch {}/{}'.format(epoch, num_epochs - 1)) 125 | print('-' * 10) 126 | 127 | # Each epoch has a training and validation phase 128 | # Due to the GPU memory limitations, we don't use the 'val' dataset 129 | for phase in ['train']: 130 | if phase == 'train': 131 | scheduler.step() 132 | model.train(True) # Set model to training mode 133 | else: 134 | model.train(False) # Set model to evaluate mode 135 | 136 | running_loss = 0.0 137 | running_corrects = 0.0 138 | # Iterate over data. 139 | for data in dataloaders[phase]: 140 | # get the inputs 141 | inputs, labels = data 142 | 143 | if use_gpu: 144 | inputs = Variable(inputs.cuda()) 145 | labels = Variable(labels.cuda()) 146 | else: 147 | inputs, labels = Variable(inputs), Variable(labels) 148 | 149 | # zero the parameter gradients 150 | optimizer.zero_grad() 151 | 152 | # forward 153 | outputs = model(inputs) 154 | 155 | part = {} 156 | sm = nn.Softmax(dim=1) 157 | num_part = 6 158 | 159 | for i in range(num_part): 160 | part[i] = outputs[i] 161 | 162 | score = sm(part[0]) + sm(part[1]) + sm(part[2]) + sm(part[3]) + sm(part[4]) + sm(part[5]) 163 | _, preds = torch.max(score.data, 1) 164 | 165 | loss = criterion(part[0], labels) 166 | for i in range(num_part - 1): 167 | loss += criterion(part[i + 1], labels) 168 | 169 | # backward + optimize only if in training phase 170 | if phase == 'train': 171 | loss.backward() 172 | optimizer.step() 173 | 174 | running_loss += loss.item() * inputs.size(0) 175 | running_corrects += torch.sum(preds == labels.data) 176 | 177 | epoch_loss = running_loss / dataset_sizes[phase] 178 | epoch_acc = running_corrects.double() / dataset_sizes[phase] 179 | 180 | print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) 181 | log_file.write('{} epoch : {} Loss: {:.4f} Acc: {:.4f}'.format(epoch, phase, 182 | epoch_loss, epoch_acc) + '\n') 183 | 184 | # if phase == 'val' 185 | if phase == 'train': 186 | last_model_wts = model.state_dict() 187 | if epoch % 10 == 9: 188 | save_network(model, epoch, stage) 189 | if epoch_acc > best_acc: 190 | best_acc = epoch_acc 191 | best_model_wts = model.state_dict() 192 | print() 193 | 194 | time_elapsed = time.time() - since 195 | print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed % 60)) 196 | print('Best val acc: {:4f}'.format(best_acc)) 197 | model.load_state_dict(best_model_wts) 198 | save_network(model, 'best', stage) 199 | model.load_state_dict(last_model_wts) 200 | save_network(model, 'last', stage) 201 | return model 202 | 203 | 204 | ###################################################################### 205 | # Save model 206 | # --------------------------- 207 | def save_network(network, epoch_label, stage): 208 | save_filename = 'net_%s.pth' % epoch_label 209 | save_path = os.path.join(save_dir, stage) 210 | if not os.path.isdir(save_path): 211 | os.mkdir(save_path) 212 | if is_parallel_train: 213 | torch.save(network.module.state_dict(), save_path + '/' + save_filename) 214 | else: 215 | torch.save(network.cpu().state_dict(), save_path + '/' + save_filename) 216 | if torch.cuda.is_available(): 217 | network.cuda(gpu_ids[0]) 218 | 219 | ###################################################################### 220 | # load model 221 | # --------------------------- 222 | def load_network(network, stage): 223 | save_path = os.path.join(args.save_dir, stage, 'net_last.pth') 224 | network.load_state_dict(torch.load(save_path)) 225 | return network 226 | 227 | ###################################################################### 228 | # PCB train 229 | # ------------------ 230 | # Step1 : train the PCB model 231 | # According to original paper, we set the difference learning rate for difference layers. 232 | 233 | 234 | def pcb_train(model, criterion, log_file, stage, num_epoch): 235 | ignored_params = list(map(id, get_net(is_parallel_train, model).classifiers.parameters())) 236 | 237 | base_params = filter(lambda p: id(p) not in ignored_params, get_net(is_parallel_train, model).parameters()) 238 | optimizer_ft = optim.SGD([ 239 | {'params': base_params, 'lr': 0.01}, 240 | {'params': get_net(is_parallel_train, model).classifiers.parameters(), 'lr': 0.1}, 241 | ], weight_decay=5e-4, momentum=0.9, nesterov=True) 242 | 243 | # Decay LR by a factor of 0.1 every 40 epochs 244 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=40, gamma=0.1) 245 | model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 246 | log_file, stage, num_epochs=num_epoch) 247 | return model 248 | 249 | 250 | ###################################################################### 251 | # RPP train 252 | # ------------------ 253 | # Setp 2&3: train the rpp layers 254 | # According to original paper, we set the learning rate at 0.01 for rpp layers. 255 | 256 | def rpp_train(model, criterion, log_file, stage, num_epoch): 257 | 258 | # ignored_params = list(map(id, get_net(opt, model).avgpool.parameters())) 259 | # base_params = filter(lambda p: id(p) not in ignored_params, get_net(opt, model).parameters()) 260 | # optimizer_ft = optim.SGD([ 261 | # {'params': base_params, 'lr': 0.00}, 262 | # {'params': get_net(opt, model).avgpool.parameters(), 'lr': 0.01}, 263 | # ], weight_decay=5e-4, momentum=0.9, nesterov=True) 264 | optimizer_ft = optim.SGD(get_net(is_parallel_train, model).avgpool.parameters(), lr=0.01, 265 | weight_decay=5e-4, momentum=0.9, nesterov=True) 266 | 267 | # Decay LR by a factor of 0.1 every 100 epochs (never use) 268 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=100, gamma=0.1) 269 | model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 270 | log_file, stage, num_epochs=num_epoch) 271 | return model 272 | 273 | ###################################################################### 274 | # full train 275 | # ------------------ 276 | # Step 4: train the whole net 277 | # According to original paper, we set the difference learning rate for the whole net 278 | 279 | def full_train(model, criterion, log_file, stage, num_epoch): 280 | ignored_params = list(map(id, get_net(is_parallel_train, model).classifiers.parameters())) 281 | ignored_params += list(map(id, get_net(is_parallel_train, model).avgpool.parameters())) 282 | 283 | base_params = filter(lambda p: id(p) not in ignored_params, get_net(is_parallel_train, model).parameters()) 284 | optimizer_ft = optim.SGD([ 285 | {'params': base_params, 'lr': 0.001}, 286 | {'params': get_net(is_parallel_train, model).classifiers.parameters(), 'lr': 0.01}, 287 | {'params': get_net(is_parallel_train, model).avgpool.parameters(), 'lr': 0.01}, 288 | ], weight_decay=5e-4, momentum=0.9, nesterov=True) 289 | 290 | # Decay LR by a factor of 0.1 every 100 epochs (never use) 291 | exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=100, gamma=0.1) 292 | model = train_model(model, criterion, optimizer_ft, exp_lr_scheduler, 293 | log_file, stage, num_epochs=num_epoch) 294 | return model 295 | 296 | 297 | ###################################################################### 298 | # Start training 299 | # --------------------------- 300 | 301 | 302 | # Record the training information # 303 | f = open(args.save_dir + 'train_log.txt', 'w') 304 | 305 | copyfile('./train.py', args.save_dir + '/train.py') 306 | copyfile('./model.py', args.save_dir + '/model.py') 307 | 308 | # step1: PCB training # 309 | stage = 'PCB' 310 | model = PCB(len(class_names)) 311 | if use_gpu: 312 | model = model.cuda() 313 | if is_parallel_train: 314 | model = nn.DataParallel(model, device_ids=gpu_ids) 315 | 316 | criterion = nn.CrossEntropyLoss() 317 | 318 | model = pcb_train(model, criterion, f, stage, 60) 319 | 320 | ############################ 321 | # step2&3: RPP training # 322 | if args.RPP: 323 | stage = 'RPP' 324 | model = get_net(is_parallel_train, model).convert_to_rpp() 325 | if use_gpu: 326 | model = model.cuda() 327 | if is_parallel_train: 328 | model = nn.DataParallel(model, device_ids=gpu_ids) 329 | 330 | model = rpp_train(model, criterion, f, stage, 5) 331 | 332 | ############################ 333 | # step4: whole net training # 334 | stage = 'full' 335 | 336 | full_train(model, criterion, f, stage, 10) 337 | f.close() 338 | 339 | 340 | 341 | 342 | 343 | 344 | --------------------------------------------------------------------------------