class DarkflowCkptClassifier:
def __init__(self,yml_config,threshold):
self.cfg_obj = yolocfg.read_yaml_main_config(yml_config)
self.options = yolocfg.generate_darkflow_args(self.cfg_obj)
self.options['threshold'] = threshold
self.options['load'] = 'best'#-1 # picks the last one
self.tfnet = TFNet(self.options)
assert self.tfnet.FLAGS.threshold>=0
def classify(self,imgcv):
if isinstance(imgcv,np.ndarray):
return self.tfnet.return_predict(imgcv)
else:
return self.tfnet.return_predict_batch(imgcv)
def classify2(self,imgcv,generate_img=True,generate_boxes=True):
assert isinstance(imgcv, np.ndarray), \
'Image is not a np.ndarray'
h, w, _ = imgcv.shape
im = self.tfnet.framework.resize_input(imgcv)
this_inp = np.expand_dims(im, 0)
feed_dict = {self.tfnet.inp : this_inp}
out = self.tfnet.sess.run(self.tfnet.out, feed_dict)
out = out.reshape(out.shape[1],out.shape[2],out.shape[3])#out[1],out[2],out[3])
newim = None
if generate_img:
outcopy = out.copy()
newim = self.tfnet.framework.postprocess(outcopy,imgcv,save=False,putText=False)
boxesInfo = list()
if generate_boxes:
boxes = self.tfnet.framework.findboxes(out)
threshold = self.tfnet.FLAGS.threshold
for box in boxes:
tmpBox = self.tfnet.framework.process_box(box, h, w, threshold)
if tmpBox is None:
continue
boxesInfo.append({
"label": tmpBox[4],
"confidence": tmpBox[6],
"topleft": {
"x": tmpBox[0],
"y": tmpBox[2]},
"bottomright": {
"x": tmpBox[1],
"y": tmpBox[3]}
})
return (newim,boxesInfo)
