import cv2
import numpy as np
import torch
from torch.autograd import Variable
import torch.nn.functional as F
import torch.nn as nn
import pickle
import os
from torch.nn.modules import CrossMapLRN2d as SpatialCrossMapLRN
from torch.autograd import Function, Variable
from torch.nn import Module
def clip_boxes(boxes, im_shape):
"""
Clip boxes to image boundaries.
"""
boxes = np.asarray(boxes)
if boxes.shape[0] == 0:
return boxes
boxes = np.copy(boxes)
boxes[:, 0::4] = np.maximum(np.minimum(boxes[:, 0::4], im_shape[1] - 1), 0)
boxes[:, 1::4] = np.maximum(np.minimum(boxes[:, 1::4], im_shape[0] - 1), 0)
boxes[:, 2::4] = np.maximum(np.minimum(boxes[:, 2::4], im_shape[1] - 1), 0)
boxes[:, 3::4] = np.maximum(np.minimum(boxes[:, 3::4], im_shape[0] - 1), 0)
return boxes
def load_net(fname, net, prefix='', load_state_dict=False):
import h5py
with h5py.File(fname, mode='r') as h5f:
h5f_is_module = True
for k in h5f.keys():
if not str(k).startswith('module.'):
h5f_is_module = False
break
if prefix == '' and not isinstance(net, nn.DataParallel) and h5f_is_module:
prefix = 'module.'
for k, v in net.state_dict().items():
k = prefix + k
if k in h5f:
param = torch.from_numpy(np.asarray(h5f[k]))
if v.size() != param.size():
print('Inconsistent shape: {}, {}'.format(v.size(), param.size()))
else:
v.copy_(param)
else:
print.warning('No layer: {}'.format(k))
epoch = h5f.attrs['epoch'] if 'epoch' in h5f.attrs else -1
if not load_state_dict:
if 'learning_rates' in h5f.attrs:
lr = h5f.attrs['learning_rates']
else:
lr = h5f.attrs.get('lr', -1)
lr = np.asarray([lr] if lr > 0 else [], dtype=np.float)
return epoch, lr
state_file = fname + '.optimizer_state.pk'
if os.path.isfile(state_file):
with open(state_file, 'rb') as f:
state_dicts = pickle.load(f)
if not isinstance(state_dicts, list):
state_dicts = [state_dicts]
else:
state_dicts = None
return epoch, state_dicts
class Inception(nn.Module):
def __init__(self, in_planes, n1x1, n3x3red, n3x3, n5x5red, n5x5, pool_planes):
super(Inception, self).__init__()
self.b1 = nn.Sequential(
nn.Conv2d(in_planes, n1x1, kernel_size=1),
nn.ReLU(True),
)
self.b2 = nn.Sequential(
nn.Conv2d(in_planes, n3x3red, kernel_size=1),
nn.ReLU(True),
nn.Conv2d(n3x3red, n3x3, kernel_size=3, padding=1),
nn.ReLU(True),
)
self.b3 = nn.Sequential(
nn.Conv2d(in_planes, n5x5red, kernel_size=1),
nn.ReLU(True),
nn.Conv2d(n5x5red, n5x5, kernel_size=5, padding=2),
nn.ReLU(True),
)
self.b4 = nn.Sequential(
nn.MaxPool2d(3, stride=1, padding=1),
nn.Conv2d(in_planes, pool_planes, kernel_size=1),
nn.ReLU(True),
)
def forward(self, x):
y1 = self.b1(x)
y2 = self.b2(x)
y3 = self.b3(x)
y4 = self.b4(x)
return torch.cat([y1,y2,y3,y4], 1)
class GoogLeNet(nn.Module):
output_channels = 832
def __init__(self):
super(GoogLeNet, self).__init__()
self.pre_layers = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
nn.ReLU(True),
nn.MaxPool2d(3, stride=2, ceil_mode=True),
SpatialCrossMapLRN(5),
nn.Conv2d(64, 64, 1),
nn.ReLU(True),
nn.Conv2d(64, 192, 3, padding=1),
nn.ReLU(True),
SpatialCrossMapLRN(5),
nn.MaxPool2d(3, stride=2, ceil_mode=True),
)
self.a3 = Inception(192, 64, 96, 128, 16, 32, 32)
self.b3 = Inception(256, 128, 128, 192, 32, 96, 64)
self.maxpool = nn.MaxPool2d(3, stride=2, ceil_mode=True)
self.a4 = Inception(480, 192, 96, 208, 16, 48, 64)
self.b4 = Inception(512, 160, 112, 224, 24, 64, 64)
self.c4 = Inception(512, 128, 128, 256, 24, 64, 64)
self.d4 = Inception(512, 112, 144, 288, 32, 64, 64)
self.e4 = Inception(528, 256, 160, 320, 32, 128, 128)
def forward(self, x):
out = self.pre_layers(x)
out = self.a3(out)
out = self.b3(out)
out = self.maxpool(out)
out = self.a4(out)
out = self.b4(out)
out = self.c4(out)
out = self.d4(out)
out = self.e4(out)
return out
class Model(nn.Module):
def __init__(self, n_parts=8):
super(Model, self).__init__()
self.n_parts = n_parts
self.feat_conv = GoogLeNet()
self.conv_input_feat = nn.Conv2d(self.feat_conv.output_channels, 512, 1)
self.conv_att = nn.Conv2d(512, self.n_parts, 1)
for i in range(self.n_parts):
setattr(self, 'linear_feature{}'.format(i+1), nn.Linear(512, 64))
def forward(self, x):
feature = self.feat_conv(x)
feature = self.conv_input_feat(feature)
att_weights = torch.sigmoid(self.conv_att(feature))
linear_feautres = []
for i in range(self.n_parts):
masked_feature = feature * torch.unsqueeze(att_weights[:, i], 1)
pooled_feature = F.avg_pool2d(masked_feature, masked_feature.size()[2:4])
linear_feautres.append(
getattr(self, 'linear_feature{}'.format(i+1))(pooled_feature.view(pooled_feature.size(0), -1))
)
concat_features = torch.cat(linear_feautres, 1)
normed_feature = concat_features / torch.clamp(torch.norm(concat_features, 2, 1, keepdim=True), min=1e-6)
return normed_feature
def load_reid_model(ckpt):
model = Model(n_parts=8)
model.inp_size = (80, 160)
load_net(ckpt, model)
print('Load ReID model from {}'.format(ckpt))
model = model.cuda()
model.eval()
return model
def im_preprocess(image):
image = np.asarray(image, np.float32)
image -= np.array([104, 117, 123], dtype=np.float32).reshape(1, 1, -1)
image = image.transpose((2, 0, 1))
return image
def extract_image_patches(image, bboxes):
bboxes = np.round(bboxes).astype(np.int)
bboxes = clip_boxes(bboxes, image.shape)
patches = [image[box[1]:box[3], box[0]:box[2]] for box in bboxes]
return patches
def extract_reid_features(reid_model, image, tlbrs):
if len(tlbrs) == 0:
return torch.FloatTensor()
patches = extract_image_patches(image, tlbrs)
patches = np.asarray([im_preprocess(cv2.resize(p, reid_model.inp_size)) for p in patches], dtype=np.float32)
with torch.no_grad():
im_var = Variable(torch.from_numpy(patches))
im_var = im_var.cuda()
features = reid_model(im_var).data
return features
