罗浩表征学习代码实践(个人学习笔记)
1.数据集(Market1501)
2.主要代码
2.1data_manager.py
from __future__ import print_function, absolute_import
import os
import os.path as osp
import glob
import re
from reid_tutorial.util.utils import mkdir_if_missing, write_json, read_json
from IPython import embed
class Market1501(object):
"""
Market1501
Reference:
Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.
URL: http://www.liangzheng.org/Project/project_reid.html
Dataset statistics:
# identities: 1501 (+1 for background)
# images: 12936 (train) + 3368 (query) + 15913 (gallery)
"""
dataset_dir = 'market1501'
def __init__(self, root='data', **kwargs):
self.dataset_dir = osp.join(root, self.dataset_dir)
self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
self.query_dir = osp.join(self.dataset_dir, 'query')
self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')
self._check_before_run()
train, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, relabel=True)
query, num_query_pids, num_query_imgs = self._process_dir(self.query_dir, relabel=False)
gallery, num_gallery_pids, num_gallery_imgs = self._process_dir(self.gallery_dir, relabel=False)
num_total_pids = num_train_pids + num_query_pids
num_total_imgs = num_train_imgs + num_query_imgs + num_gallery_imgs
print("=> Market1501 loaded")
print("Dataset statistics:")
print(" ------------------------------")
print(" subset | # ids | # images")
print(" ------------------------------")
print(" train | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
print(" query | {:5d} | {:8d}".format(num_query_pids, num_query_imgs))
print(" gallery | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_imgs))
print(" ------------------------------")
print(" total | {:5d} | {:8d}".format(num_total_pids, num_total_imgs))
print(" ------------------------------")
self.train = train
self.query = query
self.gallery = gallery
self.num_train_pids = num_train_pids
self.num_query_pids = num_query_pids
self.num_gallery_pids = num_gallery_pids
def _check_before_run(self): #判断文件夹是否存在
if not os.path.exists(self.dataset_dir):
raise RuntimeError("'{}' is not available".format(self.dataset_dir))
if not osp.exists(self.train_dir):
raise RuntimeError("'{}' is not available".format(self.train_dir))
if not osp.exists(self.query_dir):
raise RuntimeError("'{}' is not available".format(self.query_dir))
if not osp.exists(self.gallery_dir):
raise RuntimeError("'{}' is not available".format(self.gallery_dir))
def _process_dir(self, dir_path, relabel=False):
img_paths = glob.glob(osp.join(dir_path, '*.jpg'))
pattern = re.compile(r'([-\d]+)_c(\d)')
pid_container = set() # 将ID集合(去重)
for img_path in img_paths:
pid, _ = map(int, pattern.search(img_path).groups())
if pid == -1: continue # junk images are just ignored
pid_container.add(pid)
pid2label = {pid:label for label, pid in enumerate(pid_container)} #形成映射关系(ID:张数)
dataset = []
for img_path in img_paths:
pid, camid = map(int, pattern.search(img_path).groups())
if pid == -1: continue # junk images are just ignored
assert 0 <= pid <= 1501 # 判断ID是否在0-1500范围之内
assert 1 <= camid <= 6
camid -= 1 # index starts from 0
if relabel: pid = pid2label[pid]
dataset.append((img_path, pid, camid))
num_pids = len(pid_container)
num_imgs = len(dataset)
return dataset, num_pids, num_imgs
"""Create dataset"""
__img_factory = {
'market1501': Market1501,
}
def get_names():
return __img_factory.keys()
def init_img_dataset(name, **kwargs):
if name not in __img_factory.keys():
raise KeyError("Invalid dataset, got '{}', but expected to be one of {}".format(name, __img_factory.keys()))
return __img_factory[name](**kwargs)
if __name__=='__main__' :
data = Market1501(root='D:/xiangmu/reid/data')
2.2dataset_loader.py
from __future__ import print_function, absolute_import
import os
from PIL import Image
from IPython import embed
import numpy as np
import os.path as osp
import torch
from torch.utils.data import Dataset
from reid_tutorial.util import data_manager
# 读取图片
def read_image(img_path):
got_img = False
if not osp.exists(img_path):
raise IOError("{} does not exist".format(img_path))
while not got_img:
try:
img = Image.open(img_path).convert('RGB')
got_img = True
except IOError:
print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path))
pass
return img
class ImageDataset(Dataset):
def __init__(self, dataset, transform=None):
self.dataset = dataset
self.transform = transform
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
img_path, pid, camid = self.dataset[index]
img = read_image(img_path)
if self.transform is not None:
img = self.transform(img)
return img, pid, camid
if __name__ == '__main__':
dataset = data_manager.init_img_dataset(root='D:/xiangmu/reid/data', name='market1501')
train_loader = ImageDataset(dataset.train)
embed()
2.3transform.py
from __future__ import absolute_import
from torchvision.transforms import *
from PIL import Image
import random
import numpy as np
class Random2DTranslation(object):
"""
With a probability, first increase image size to (1 + 1/8), and then perform random crop.
Args:
height (int): target height.
width (int): target width.
p (float): probability of performing this transformation. Default: 0.5.
"""
def __init__(self, height, width, p=0.5, interpolation=Image.BILINEAR):
self.height = height
self.width = width
self.p = p
self.interpolation = interpolation
def __call__(self, img):
"""
Args:
img (PIL Image): Image to be cropped.
Returns:
PIL Image: Cropped image.
"""
if random.random() < self.p:
return img.resize((self.width, self.height), self.interpolation)
new_width, new_height = int(round(self.width * 1.125)), int(round(self.height * 1.125))
resized_img = img.resize((new_width, new_height), self.interpolation)
x_maxrange = new_width - self.width
y_maxrange = new_height - self.height
x1 = int(round(random.uniform(0, x_maxrange)))
y1 = int(round(random.uniform(0, y_maxrange)))
croped_img = resized_img.crop((x1, y1, x1 + self.width, y1 + self.height))
return croped_img
if __name__ == '__main__':
img = Image.open('D:/xiangmu/reid/data/market1501/bounding_box_train/0002_c1s1_000801_01.jpg')#打开一张图像
#transform = Random2DTranslation(256, 128, 0.5)
import torchvision.transforms as transforms
transform = transforms.Compose(
[
Random2DTranslation(256, 128, 0.5),
transforms.RandomHorizontalFlip(),#水平翻转
transforms.ToTensor(),
]
)
img_t = transform(img)#裁剪后的图像
print(img_t.shape)
# import matplotlib.pyplot as plt
#
# plt.figure(12)
# plt.subplot(121)
# plt.imshow(img)
# plt.subplot(122)
# plt.imshow(img_t)
# plt.show()
2.4ResNet.py
from __future__ import absolute_import
import torch
from torch import nn
from torch.nn import functional as F
import torchvision
from IPython import embed
class ResNet50(nn.Module):
def __init__(self, num_classes, loss={'xent'}, **kwargs):
super(ResNet50, self).__init__()
self.loss = loss
resnet50 = torchvision.models.resnet50(pretrained=True)
self.base = nn.Sequential(*list(resnet50.children())[:-2])#去掉分类器层
self.classifier = nn.Linear(2048, num_classes)#定义自己的分类器
def forward(self, x):
x = self.base(x)
x = F.avg_pool2d(x, x.size()[2:])
f = x.view(x.size(0), -1)#展平
# f = 1.*f / (torch.norm(f, 2, dim=-1, keepdim=True).expand_as(f) + 1e-12)
if not self.training:
return f
y = self.classifier(f)
return y
if __name__ == '__main__':
model = ResNet50(num_classes=751)
imgs = torch.Tensor(32, 3, 256, 128)
f = model(imgs)
2.5train_class.py
from __future__ import absolute_import
import os
import sys
import time
import datetime
import argparse
import os.path as osp
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
from torch.optim import lr_scheduler
import models
from util.losses import CrossEntropyLoss, DeepSupervision, CrossEntropyLabelSmooth, TripletLossAlignedReID
from util import data_manager
from util import transforms as T
from util.dataset_loader import ImageDataset
from util.utils import Logger
from util.utils import AverageMeter, Logger, save_checkpoint
from util.eval_metrics import evaluate
from util.optimizers import init_optim
from util.samplers import RandomIdentitySampler
from IPython import embed
parser = argparse.ArgumentParser(description='Train AlignedReID with cross entropy loss and triplet hard loss')
# Datasets
parser.add_argument('--root', type=str, default='D:/xiangmu/reid/data', help="root path to data directory")
parser.add_argument('-d', '--dataset', type=str, default='market1501',
choices=data_manager.get_names())
parser.add_argument('-j', '--workers', default=4, type=int,
help="number of data loading workers (default: 4)")
parser.add_argument('--height', type=int, default=256,
help="height of an image (default: 256)")
parser.add_argument('--width', type=int, default=128,
help="width of an image (default: 128)")
parser.add_argument('--split-id', type=int, default=0, help="split index")
# CUHK03-specific setting
parser.add_argument('--cuhk03-labeled', action='store_true',
help="whether to use labeled images, if false, detected images are used (default: False)")
parser.add_argument('--cuhk03-classic-split', action='store_true',
help="whether to use classic split by Li et al. CVPR'14 (default: False)")
parser.add_argument('--use-metric-cuhk03', action='store_true',
help="whether to use cuhk03-metric (default: False)")
# Optimization options
parser.add_argument('--labelsmooth', action='store_true', help="label smooth")
parser.add_argument('--optim', type=str, default='adam', help="optimization algorithm (see optimizers.py)")
parser.add_argument('--max-epoch', default=1, type=int,
help="maximum epochs to run")
parser.add_argument('--start-epoch', default=0, type=int,
help="manual epoch number (useful on restarts)")
parser.add_argument('--train-batch', default=15, type=int,
help="train batch size")
parser.add_argument('--test-batch', default=15, type=int, help="test batch size")
parser.add_argument('--lr', '--learning-rate', default=0.0002, type=float,
help="initial learning rate")
parser.add_argument('--stepsize', default=20, type=int,
help="stepsize to decay learning rate (>0 means this is enabled)")
parser.add_argument('--gamma', default=0.1, type=float,
help="learning rate decay")
parser.add_argument('--weight-decay', default=5e-04, type=float, help="weight decay (default: 5e-04)")
# Architecture
parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.get_names())
# Miscs
parser.add_argument('--print-freq', type=int, default=10, help="print frequency")
parser.add_argument('--seed', type=int, default=1, help="manual seed")
parser.add_argument('--resume', type=str, default='', metavar='PATH')
parser.add_argument('--evaluate', action='store_true', help="evaluation only")
parser.add_argument('--eval-step', type=int, default=-1,
help="run evaluation for every N epochs (set to -1 to test after training)")
parser.add_argument('--start-eval', type=int, default=0, help="start to evaluate after specific epoch")
parser.add_argument('--save-dir', type=str, default='log')
parser.add_argument('--use_cpu', action='store_true', help="use cpu")
parser.add_argument('--gpu-devices', default='0', type=str, help='gpu device ids for CUDA_VISIBLE_DEVICES')
parser.add_argument('--reranking', action='store_true', help='result re_ranking')
parser.add_argument('--test_distance', type=str, default='global', help='test distance type')
parser.add_argument('--unaligned', action='store_true', help='test local feature with unalignment')
args = parser.parse_args()
def test(model, queryloader, galleryloader, use_gpu, ranks=[1, 5, 10, 20]):
batch_time = AverageMeter()
model.eval()
with torch.no_grad():
qf, q_pids, q_camids, lqf = [], [], [], []
for batch_idx, (imgs, pids, camids) in enumerate(queryloader):
if use_gpu: imgs = imgs.cuda()
end = time.time()
features = model(imgs)
local_features = model(imgs)
batch_time.update(time.time() - end)
features = features.data.cpu()
local_features = local_features.data.cpu()
qf.append(features)
lqf.append(local_features)
q_pids.extend(pids)
q_camids.extend(camids)
qf = torch.cat(qf, 0)
lqf = torch.cat(lqf, 0)
q_pids = np.asarray(q_pids)
q_camids = np.asarray(q_camids)
print("Extracted features for query set, obtained {}-by-{} matrix".format(qf.size(0), qf.size(1)))
gf, g_pids, g_camids, lgf = [], [], [], []
end = time.time()
for batch_idx, (imgs, pids, camids) in enumerate(galleryloader):
if use_gpu: imgs = imgs.cuda()
end = time.time()
features = model(imgs)
local_features = model(imgs)
batch_time.update(time.time() - end)
features = features.data.cpu()
local_features = local_features.data.cpu()
gf.append(features)
lgf.append(local_features)
g_pids.extend(pids)
g_camids.extend(camids)
gf = torch.cat(gf, 0)
lgf = torch.cat(lgf, 0)
g_pids = np.asarray(g_pids)
g_camids = np.asarray(g_camids)
print("Extracted features for gallery set, obtained {}-by-{} matrix".format(gf.size(0), gf.size(1)))
print("==> BatchTime(s)/BatchSize(img): {:.3f}/{}".format(batch_time.avg, args.test_batch))
# feature normlization
qf = 1. * qf / (torch.norm(qf, 2, dim=-1, keepdim=True).expand_as(qf) + 1e-12)
gf = 1. * gf / (torch.norm(gf, 2, dim=-1, keepdim=True).expand_as(gf) + 1e-12)
m, n = qf.size(0), gf.size(0)
distmat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
#distmat.addmm_(1, -2, qf, gf.t())
distmat.addmm_(qf, gf.t(), beta=1, alpha=-2)
distmat = distmat.numpy()
if not args.test_distance == 'global':
print("Only using global branch")
from reid_tutorial import low_memory_local_dist
lqf = lqf.permute(0, 2, 1)
lgf = lgf.permute(0, 2, 1)
local_distmat = low_memory_local_dist(lqf.numpy(), lgf.numpy(), aligned=not args.unaligned)
if args.test_distance == 'local':
print("Only using local branch")
distmat = local_distmat
if args.test_distance == 'global_local':
print("Using global and local branches")
distmat = local_distmat+distmat
print("Computing CMC and mAP")
cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, use_metric_cuhk03=args.use_metric_cuhk03)
print("Results ----------")
print("mAP: {:.1%}".format(mAP))
print("CMC curve")
for r in ranks:
print("Rank-{:<3}: {:.1%}".format(r, cmc[r - 1]))
print("------------------")
if args.reranking:
from reid_tutorial import re_ranking
if args.test_distance == 'global':
print("Only using global branch for reranking")
distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3)
else:
local_qq_distmat = low_memory_local_dist(lqf.numpy(), lqf.numpy(), aligned=not args.unaligned)
local_gg_distmat = low_memory_local_dist(lgf.numpy(), lgf.numpy(), aligned=not args.unaligned)
local_dist = np.concatenate(
[np.concatenate([local_qq_distmat, local_distmat], axis=1),
np.concatenate([local_distmat.T, local_gg_distmat], axis=1)],
axis=0)
if args.test_distance == 'local':
print("Only using local branch for reranking")
distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3, local_distmat=local_dist, only_local=True)
elif args.test_distance == 'global_local':
print("Using global and local branches for reranking")
distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3, local_distmat=local_dist, only_local=False)
print("Computing CMC and mAP for re_ranking")
cmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, use_metric_cuhk03=args.use_metric_cuhk03)
print("Results ----------")
print("mAP(RK): {:.1%}".format(mAP))
print("CMC curve(RK)")
for r in ranks:
print("Rank-{:<3}: {:.1%}".format(r, cmc[r - 1]))
print("------------------")
return cmc[0]
def main():
use_gpu = torch.cuda.is_available()
if args.use_cpu:
use_gpu = False
if use_gpu:
pin_memory = True
else:
pin_memory = False
if not args.evaluate:
sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt'))
else:
sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt'))
print("==========\nArgs:{}\n==========".format(args))
if use_gpu:
print("Currently using GPU {}".format(args.gpu_devices))
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices
cudnn.benchmark = True
torch.cuda.manual_seed_all(args.seed)
else:
print("Currently using CPU (GPU is highly recommended)")
print("Initializing dataset {}".format(args.dataset))
dataset = data_manager.init_img_dataset(
root=args.root, name=args.dataset, split_id=args.split_id,)
# data augmentation
transform_train = T.Compose([
T.Random2DTranslation(args.height, args.width),
T.RandomHorizontalFlip(),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = T.Compose([
T.Resize((args.height, args.width)),
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
trainloader = DataLoader(
ImageDataset(dataset.train, transform=transform_train),
#sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances),
batch_size=args.train_batch, num_workers=args.workers,
pin_memory=pin_memory, drop_last=True,
)
queryloader = DataLoader(
ImageDataset(dataset.query, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
galleryloader = DataLoader(
ImageDataset(dataset.gallery, transform=transform_test),
batch_size=args.test_batch, shuffle=False, num_workers=args.workers,
pin_memory=pin_memory, drop_last=False,
)
print("Initializing model: {}".format(args.arch))
model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss='softmax')
print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0))
criterion_class = CrossEntropyLoss()
# criterion_class = nn.CrossEntropyLoss()
# optimizer = torch.optim.adam()
# 优化器
optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)
# 学习率
if args.stepsize > 0:
scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) # 阶梯型衰减
start_epoch = args.start_epoch
if args.resume:
print("Loading checkpoint from '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint['state_dict'])
start_epoch = checkpoint['epoch']
if use_gpu:
model = nn.DataParallel(model).cuda()
# model.module.parameter()
if args.evaluate:
print('Evaluate only!')
test(model, queryloader, galleryloader, use_gpu)
return 0
start_time = time.time()
train_time = 0
best_rank1 = -np.inf
best_epoch = 0
print('start training')
for epoch in range(start_epoch, args.max_epoch):
start_train_time = time.time()
train(epoch, model, criterion_class, optimizer, trainloader, use_gpu)
#model.save_state_dict()
train_time += round(time.time() - start_train_time)
if args.stepsize > 0: scheduler.step()
if (epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch:
print("==> Test")
rank1 = test(model, queryloader, galleryloader, use_gpu)
is_best = rank1 > best_rank1
if is_best:
best_rank1 = rank1
best_epoch = epoch + 1
if use_gpu:
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
save_checkpoint({
'state_dict': state_dict,
'rank1': rank1,
'epoch': epoch,
}, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar'))
print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch))
elapsed = round(time.time() - start_time)
elapsed = str(datetime.timedelta(seconds=elapsed))
train_time = str(datetime.timedelta(seconds=train_time))
print("Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.".format(elapsed, train_time))
def train(epoch, model, criterion_class, optimizer, trainloader, use_gpu):
model.train()
losses = AverageMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
end = time.time()
for batch_idx, (imgs, pids, _) in enumerate(trainloader):
if use_gpu: # GPU跑数据
imgs, pids = imgs.cuda(), pids.cuda()
# measure data loading time
data_time.update(time.time() - end)
outputs = model(imgs)
loss = criterion_class(outputs, pids)
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
losses.update(loss.item(), pids.size(0))
if (batch_idx+1) % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
epoch+1, batch_idx+1, len(trainloader), batch_time=batch_time, data_time=data_time, loss=losses))
if __name__ == '__main__':
main()
3.具体实现的项目的地址如下:
GitHub - yizhanchengshou/biaozhengxuexiContribute to yizhanchengshou/biaozhengxuexi development by creating an account on GitHub.
https://github.com/yizhanchengshou/biaozhengxuexi.git
4.最终训练出的模型:
最终的效果不好,还需继续进行参数的调整。