def segmentation(model, input_folder, output_folder, scale):
keras_weights_file = model
print('start processing...')
# load model
model = get_testing_model_resnet101()
model.load_weights(keras_weights_file)
params, model_params = config_reader()
scale_list = []
for item in scale:
scale_list.append(float(item))
params['scale_search'] = scale_list
seg_dict = {}
kpts_dict = {}
# generate image with body parts
for filename in os.listdir(input_folder):
if filename.endswith(".png") or filename.endswith(".jpg"):
print(input_folder + '/' + filename)
#------------------This is what you need------------------------------------------------
#kpts should contain what you need
canvas, seg, kpts = process(input_folder + '/' + filename, params,
model_params, model)
#specifically, it is a dictionary with keys 1, 2, and 5 (rather arbitrary for now)
assert 1 in kpts.keys()
assert 2 in kpts.keys()
assert 5 in kpts.keys()
#kpts[1] should be a tuple of neck coords, kpts[2] left shoulder, and kpts[5] right shoulder
#Use them in cv2 order, which is to say the tuples should be ordered (ycoord, xcoord)
# ------------------This is what you need------------------------------------------------
cv2.imwrite(output_folder + '/sk_' + filename, canvas)
seg_argmax = np.argmax(seg, axis=-1)
seg_max = np.max(seg, axis=-1)
seg_max_thres = (seg_max > 0.1).astype(np.uint8)
seg_argmax *= seg_max_thres
seg_dict[filename] = seg_argmax
kpts_dict[filename] = kpts
#not completely necessary
filename = '%s/%s.jpg' % (output_folder,
'seg_' + os.path.splitext(filename)[0])
cv2.imwrite(filename, seg_argmax)
return seg_dict, kpts_dict
candidate[subset_int[i][j], 0], candidate[subset_int[i][j],
1]
])
# code.interact(local=locals())
return canvas, heatmap_avg, paf_avg, people, seg_avg
if __name__ == '__main__':
args = parser.parse_args()
keras_weights_file = args.model
print('start processing...')
# load model
model = get_testing_model_resnet101()
model.load_weights(keras_weights_file)
params, model_params = config_reader()
scale_list = []
for item in args.scale:
scale_list.append(float(item))
params['scale_search'] = scale_list
# generate image with body parts
for filename in os.listdir(args.input_folder):
if filename.endswith(".png") or filename.endswith(".jpg"):
print(args.input_folder + '/' + filename)
canvas, heatmap, paf, people, seg = process(
args.input_folder + '/' + filename, params, model_params)