Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
258 lines (227 sloc) 10.3 KB
import os
import numpy as np
import utils.cuda_util
from keras import Input
from keras import backend as K
from keras.applications.resnet50 import preprocess_input
from keras.callbacks import EarlyStopping, ReduceLROnPlateau
from keras.engine import Model
from keras.layers import Lambda, Dense, Dropout, Flatten
from keras.models import load_model
from keras.optimizers import SGD
from keras.preprocessing import image
from keras.utils import plot_model, to_categorical
from numpy.random import randint, shuffle, choice
from baseline.train import softmax_pretrain_on_dataset
def mix_data_prepare(data_list_path, train_dir_path):
class_img_labels = dict()
class_cnt = -1
last_label = -2
last_type = ''
with open(data_list_path, 'r') as f:
for line in f:
line = line.strip()
img = line
lbl = int(line.split('_')[0])
img_type = line.split('.')[-1]
if lbl != last_label or img_type != last_type:
class_cnt = class_cnt + 1
cur_list = list()
class_img_labels[str(class_cnt)] = cur_list
last_label = lbl
last_type = img_type
img = image.load_img(os.path.join(train_dir_path, img), target_size=[224, 224])
img = image.img_to_array(img)
img = np.expand_dims(img, axis=0)
img = preprocess_input(img)
class_img_labels[str(class_cnt)].append(img[0])
return class_img_labels
def reid_data_prepare(data_list_path, train_dir_path):
if 'mix' in data_list_path:
return mix_data_prepare(data_list_path, train_dir_path)
class_img_labels = dict()
class_cnt = -1
last_label = -2
with open(data_list_path, 'r') as f:
for line in f:
line = line.strip()
img = line
lbl = int(line.split('_')[0])
if lbl != last_label:
class_cnt = class_cnt + 1
cur_list = list()
class_img_labels[str(class_cnt)] = cur_list
last_label = lbl
img = image.load_img(os.path.join(train_dir_path, img), target_size=[224, 224])
img = image.img_to_array(img)
img = np.expand_dims(img, axis=0)
img = preprocess_input(img)
class_img_labels[str(class_cnt)].append(img[0])
return class_img_labels
def pair_generator(class_img_labels, batch_size, train=False):
cur_epoch = 0
pos_prop = 5
while True:
left_label = randint(len(class_img_labels), size=batch_size)
if cur_epoch % pos_prop == 0:
right_label = left_label
else:
right_label = np.copy(left_label)
shuffle(right_label)
# select by label
left_images = list()
right_images = list()
if train:
slice_start = 0
else:
# val
slice_start = 0.9
for i in range(batch_size):
len_left_label_i = len(class_img_labels[str(left_label[i])])
left_images.append(class_img_labels[str(left_label[i])][int(slice_start * len_left_label_i):][
choice(len_left_label_i - int(len_left_label_i * slice_start))])
len_right_label_i = len(class_img_labels[str(right_label[i])])
right_images.append(class_img_labels[str(right_label[i])][int(slice_start * len_right_label_i):][
choice(len_right_label_i - int(len_right_label_i * slice_start))])
left_images = np.array(left_images)
right_images = np.array(right_images)
binary_label = (left_label == right_label).astype(int)
left_label = to_categorical(left_label, num_classes=len(class_img_labels))
right_label = to_categorical(right_label, num_classes=len(class_img_labels))
cur_epoch += 1
yield [left_images, right_images], [left_label, right_label, binary_label]
def eucl_dist(inputs):
x, y = inputs
# return K.mean(K.square((x - y)), axis=1)
return K.square((x - y))
def dis_sigmoid(dis):
return K.expand_dims(2/(1+K.exp(dis)))
def pair_model(source_model_path, num_classes):
softmax_model = load_model(source_model_path)
# base_model = ResNet50(weights='imagenet', include_top=False, input_tensor=Input(shape=(224, 224, 3)))
base_model = Model(inputs=softmax_model.input, outputs=[softmax_model.get_layer('avg_pool').output], name='resnet50')
img1 = Input(shape=(224, 224, 3), name='img_1')
img2 = Input(shape=(224, 224, 3), name='img_2')
feature1 = Flatten()(base_model(img1))
feature2 = Flatten()(base_model(img2))
dis = Lambda(eucl_dist, name='square')([feature1, feature2])
# judge = Lambda(dis_sigmoid, name='bin_out')(dis)
judge = Dense(1, activation='sigmoid', name='bin_out')(Dropout(0.9)(dis))
category_predict1 = Dense(num_classes, activation='softmax', name='ctg_out_1')(
Dropout(0.9)(feature1)
)
category_predict2 = Dense(num_classes, activation='softmax', name='ctg_out_2')(
Dropout(0.9)(feature2)
)
model = Model(inputs=[img1, img2], outputs=[category_predict1, category_predict2, judge])
model.get_layer('ctg_out_1').set_weights(softmax_model.get_layer('fc8').get_weights())
model.get_layer('ctg_out_2').set_weights(softmax_model.get_layer('fc8').get_weights())
plot_model(model, to_file='model_combined.png')
# for layer in base_model.layers[:-10]:
# layer.trainable = False
for layer in base_model.layers:
layer.trainable = True
return model
def common_lr(epoch):
if epoch < 20:
return 0.01
else:
return 0.001
def pair_tune(source_model_path, train_generator, val_generator, tune_dataset, batch_size=48, num_classes=751):
model = pair_model(source_model_path, num_classes)
model.compile(optimizer=SGD(lr=0.001, momentum=0.9),
loss={'ctg_out_1': 'categorical_crossentropy',
'ctg_out_2': 'categorical_crossentropy',
'bin_out': 'binary_crossentropy'},
loss_weights={
'ctg_out_1': 0.5,
'ctg_out_2': 0.5,
'bin_out': 1.
},
metrics=['accuracy'])
early_stopping = EarlyStopping(monitor='val_loss', patience=4)
auto_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=0, mode='auto', epsilon=0.0001,
cooldown=0, min_lr=0)
# save_model = ModelCheckpoint('resnet50-{epoch:02d}-{val_ctg_out_1_acc:.2f}.h5', period=2)
model.fit_generator(train_generator,
steps_per_epoch=16500 / batch_size + 1,
epochs=20,
validation_data=val_generator,
validation_steps=1800 / batch_size + 1,
callbacks=[auto_lr, early_stopping])
model.save(tune_dataset + '_pair_pretrain.h5')
def pair_pretrain_on_dataset(source, project_path='/home/cwh/coding/rank-reid', dataset_parent='/home/cwh/coding'):
if source == 'market':
train_list = project_path + '/dataset/market_train.list'
train_dir = dataset_parent + '/Market-1501/train'
class_count = 751
elif source == 'markets1':
train_list = project_path + '/dataset/markets1_train.list'
train_dir = dataset_parent + '/markets1'
class_count = 751
elif source == 'grid':
train_list = project_path + '/dataset/grid_train.list'
train_dir = dataset_parent + '/grid_label'
class_count = 250
elif source == 'cuhk':
train_list = project_path + '/dataset/cuhk_train.list'
train_dir = dataset_parent + '/cuhk01'
class_count = 971
elif source == 'viper':
train_list = project_path + '/dataset/viper_train.list'
train_dir = dataset_parent + '/viper'
class_count = 630
elif source == 'duke':
train_list = project_path + '/dataset/duke_train.list'
train_dir = dataset_parent + '/DukeMTMC-reID/train'
class_count = 702
elif 'grid-cv' in source:
cv_idx = int(source.split('-')[-1])
train_list = project_path + '/dataset/grid-cv/%d.list' % cv_idx
train_dir = dataset_parent + '/grid_train_probe_gallery/cross%d/train' % cv_idx
class_count = 125
elif 'mix' in source:
train_list = project_path + '/dataset/mix.list'
train_dir = dataset_parent + '/cuhk_grid_viper_mix'
class_count = 250 + 971 + 630
else:
train_list = 'unknown'
train_dir = 'unknown'
class_count = -1
class_img_labels = reid_data_prepare(train_list, train_dir)
batch_size = 16
pair_tune(
source + '_softmax_pretrain.h5',
pair_generator(class_img_labels, batch_size=batch_size, train=True),
pair_generator(class_img_labels, batch_size=batch_size, train=False),
source,
batch_size=batch_size, num_classes=class_count
)
if __name__ == '__main__':
# sources = ['cuhk_grid_viper_mix']
# sources = ['cuhk', 'viper', 'market','duke']
sources = ['duke']
for source in sources:
softmax_pretrain_on_dataset(source,
project_path='/home/cwh/coding/rank-reid',
dataset_parent='/home/cwh/coding/')
pair_pretrain_on_dataset(source)
# sources = ['grid-cv-%d' % i for i in range(10)]
# for source in sources:
# softmax_pretrain_on_dataset(source,
# project_path='/home/cwh/coding/rank-reid',
# dataset_parent='/home/cwh/coding')
# pair_pretrain_on_dataset(source,
# project_path='/home/cwh/coding/rank-reid',
# dataset_parent='/home/cwh/coding')
sources = ['grid']
for source in sources:
pair_pretrain_on_dataset(source)
# sources = ['grid-cv-%d' % i for i in range(10)]
# for source in sources:
# softmax_pretrain_on_dataset(source,
# project_path='/home/cwh/coding/rank-reid',
# dataset_parent='/home/cwh/coding')
# pair_pretrain_on_dataset(source,
# project_path='/home/cwh/coding/rank-reid',
# dataset_parent='/home/cwh/coding')