def result_write(result_file, dress_type):
print('--Parsing Config File')
params = process_config('config.cfg', dress_type)
print('--Creating Dataset')
img_dir = params['test_img_directory']
data_file = params['testing_data_file']
joint_list = params['joint_list']
dataset = DataGenerator(dress_type,
params['joint_list'],
img_dir,
test_data_file=data_file)
model = HourglassModel(nFeat=params['nfeats'],
nStack=params['nstacks'],
nModules=params['nmodules'],
nLow=params['nlow'],
outputDim=params['num_joints'],
batch_size=params['batch_size'],
attention=params['mcam'],
training=True,
drop_rate=params['dropout_rate'],
lear_rate=params['learning_rate'],
decay=params['learning_rate_decay'],
decay_step=params['decay_step'],
dataset=dataset,
name=params['name'],
logdir_train=params['log_dir_train'],
logdir_test=params['log_dir_test'],
tiny=params['tiny'],
w_loss=params['weighted_loss'],
joints=joint_list,
modif=False)
model.generate_model()
model.restore(load=params['model_file'] + params['load_file'])
print('START')
for i in range(result_file.shape[0]):
if result_file.at[
i,
'image_category'] == dress_type: # Only take the type we want
name = str(result_file.at[i, 'image_id'])
print('Dating ', i, ' PICTURE', ' name is:', name)
image = dataset.open_img(name)
output, keypoint = model.get_output(name)
heatmaps = np.zeros(
(10, output.shape[0], output.shape[1], output.shape[2]))
heatmaps[0] = output
for j in range(1, 10):
heatmaps[j], keypoint = model.get_output(name)
output = dataset.average_heatmaps(heatmaps)
keypoints = dataset.find_nkeypoints(output)
for k in range(keypoints.shape[0]):
position = str(keypoints[k][0]) + '_' + str(
keypoints[k][1]) + '_' + '1'
result_file.loc[i, joint_list[k]] = position
point = tuple(keypoints[k])
cv2.circle(image, point, 5, (0, 0, 255), -1)
cv2.imshow('image', image)
# Wait
if cv2.waitKey(10) == 27:
print('Ended')
cv2.destroyAllWindows()
break
cv2.destroyAllWindows()
model.close_sess() # 释放资源
print('FINISH')
decay=params['learning_rate_decay'],
decay_step=params['decay_step'],
dataset=dataset,
name=params['name'],
logdir_train=params['log_dir_train'],
logdir_test=params['log_dir_test'],
tiny=params['tiny'],
w_loss=params['weighted_loss'],
joints=params['joint_list'],
modif=False)
model.generate_model()
model.restore(load=params['model_file'] + params['load_file'])
test_table = dataset.read_test_data()
for i in range(len(test_table)):
image = dataset.open_img(test_table[i])
output, keypoints = model.get_output(test_table[i])
print('image_width', image.shape[1], ' image_height', image.shape[0])
print('keypoints:', end=' ')
for p in range(keypoints.shape[0]):
print(keypoints[p], end=' ')
print(' ')
for k in range(keypoints.shape[0]):
# cv2.imshow(str('point' + str(k)), output[k])
point = tuple(keypoints[k])
cv2.circle(image, point, 5,
(k / keypoints.shape[0] * 255,
(1 - k / keypoints.shape[0]) * 255, 255), -1)
output = output.astype(np.uint8)
cv2.imshow('image', image)
# Wait
