-
Notifications
You must be signed in to change notification settings - Fork 235
/
predict.py
30 lines (22 loc) · 1.02 KB
/
predict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Author: kerlomz <[email protected]>
from config import ModelConfig
def decode_maps(categories):
return {index: category for index, category in enumerate(categories, 0)}
def predict_func(image_batch, _sess, dense_decoded, op_input, model: ModelConfig, output_split=None):
output_split = model.output_split if output_split is None else output_split
category_split = model.category_split if model.category_split else ""
dense_decoded_code = _sess.run(dense_decoded, feed_dict={
op_input: image_batch,
})
decoded_expression = []
for item in dense_decoded_code:
expression = []
for i in item:
if i == -1 or i == model.category_num:
expression.append("")
else:
expression.append(decode_maps(model.category)[i])
decoded_expression.append(category_split.join(expression))
return output_split.join(decoded_expression) if len(decoded_expression) > 1 else decoded_expression[0]