def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
dataset = args.dataset # 'A' or 'B'
train_path = cfg.TRAIN_PATH.format(dataset)
train_gt_path = cfg.TRAIN_GT_PATH.format(dataset)
val_path = cfg.VAL_PATH.format(dataset)
val_gt_path = cfg.VAL_GT_PATH.format(dataset)
train_data_loader = DataLoader(train_path,
train_gt_path,
shuffle=True,
gt_downsample=False)
val_data_loader = DataLoader(val_path,
val_gt_path,
shuffle=False,
gt_downsample=False)
# 加载数据
train_X, train_Y_den, train_Y_class = train_data_loader.load_all()
val_X, val_Y_den, val_Y_class = val_data_loader.load_all()
class_weights = train_data_loader.get_class_weights()
# 定义模型 & 编译
input_shape = (None, None, 1)
model = CMTL(input_shape)
adam = Adam(lr=0.00001)
loss = {'density': 'mse', 'cls': 'categorical_crossentropy'}
loss_weights = {'density': 1.0, 'cls': 0.0001}
print('[INFO] Compiling model ...'.format(dataset))
model.compile(optimizer=adam,
loss=loss,
loss_weights=loss_weights,
metrics={
'density': [mae, mse],
'cls': 'accuracy'
})
# 定义callback
checkpointer_best_train = ModelCheckpoint(filepath=os.path.join(
cfg.MODEL_DIR, 'mcnn_' + dataset + '_train.hdf5'),
monitor='loss',
verbose=1,
save_best_only=True,
mode='min')
callback_list = [checkpointer_best_train]
# # 随机数据增广
# print('[INFO] Random data augment ...'.format(dataset))
# train_X, train_Y_den = train_data_loader.random_augment(train_X, train_Y_den)
# 训练
print('[INFO] Training Part_{} ...'.format(dataset))
model.fit(train_X, {
"density": train_Y_den,
"cls": train_Y_class
},
validation_data=(val_X, {
"density": val_Y_den,
"cls": val_Y_class
}),
batch_size=cfg.TRAIN_BATCH_SIZE,
epochs=cfg.EPOCHS,
callbacks=callback_list,
class_weight={"cls": class_weights})
def main(args):
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
dataset = args.dataset # 'A' or 'B'
if dataset == 'A':
model_path = './trained_models/mcnn_A_train.hdf5'
else:
model_path = './trained_models/mcnn_B_train.hdf5'
output_dir = './output_{}/'.format(dataset)
heatmaps_dir = os.path.join(output_dir,
'heatmaps') # directory to save heatmap
results_txt = os.path.join(output_dir,
'results.txt') # file to save predicted results
for _dir in [output_dir, heatmaps_dir]:
if not os.path.exists(_dir):
os.mkdir(_dir)
test_path = cfg.TEST_PATH.format(dataset)
test_gt_path = cfg.TEST_GT_PATH.format(dataset)
# load test set
data_loader = DataLoader(test_path,
test_gt_path,
shuffle=False,
gt_downsample=True)
_ = data_loader.load_all()
# load model
print('[INFO] Load model ...')
model = load_model(model_path, custom_objects={'tf': tf})
# test
print('[INFO] Testing Part_{} ...'.format(dataset))
mae = 0.0
mse = 0.0
acc = 0.0
for blob in data_loader:
img = blob['data']
gt_den = blob['gt_den']
gt_cls = np.argmax(blob['gt_class'])
pred_den, pred_cls = model.predict(np.expand_dims(img, axis=0))
if np.argmax(pred_cls[0]) == gt_cls:
acc += 1
gt_count = np.sum(gt_den)
pred_count = np.sum(pred_den)
mae += abs(gt_count - pred_count)
mse += ((gt_count - pred_count) * (gt_count - pred_count))
# # create and save heatmap
# pred = np.squeeze(pred) # shape(1, h, w, 1) -> shape(h, w)
# save_heatmap(pred, blob, test_path, heatmaps_dir)
# save results
with open(results_txt, 'a') as f:
line = '<{}> {:.2f}--{:.2f}\t{}--{}\n'.format(
blob['fname'].split('.')[0], gt_count, pred_count, gt_cls,
np.argmax(pred_cls[0]))
f.write(line)
mae = mae / data_loader.num_samples
mse = np.sqrt(mse / data_loader.num_samples)
acc = acc / data_loader.num_samples
print('[RESULT] MAE: %0.2f, MSE: %0.2f, Acc: %0.2f' % (mae, mse, acc))
with open(results_txt, 'a') as f:
f.write('MAE: %0.2f, MSE: %0.2f, Acc: %0.2f' % (mae, mse, acc))
