def detect_image(self, image):
if self.model_image_size != (None, None):
assert self.model_image_size[
0] % 32 == 0, 'Multiples of 32 required'
assert self.model_image_size[
1] % 32 == 0, 'Multiples of 32 required'
image_data = preprocess_image(image, self.model_image_size)
#origin image shape, in (height, width) format
image_shape = tuple(reversed(image.size))
start = time.time()
out_boxes, out_classes, out_scores = self.predict(
image_data, image_shape)
print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
end = time.time()
print("Inference time: {:.8f}s".format(end - start))
#draw result on input image
image_array = np.array(image, dtype='uint8')
image_array = draw_boxes(image_array, out_boxes, out_classes,
out_scores, self.class_names, self.colors)
return Image.fromarray(image_array)
def detect_image(self, image):
if self.model_image_size != (None, None):
assert self.model_image_size[
0] % 32 == 0, 'Multiples of 32 required'
assert self.model_image_size[
1] % 32 == 0, 'Multiples of 32 required'
image = Image.open(image)
image_data = preprocess_image(image, self.model_image_size)
image_shape = image.size
self.interpreter.set_tensor(self.input_details[0]['index'], image_data)
start = time.time()
self.interpreter.invoke()
output_data = [
self.interpreter.get_tensor(self.output_details[2]['index']),
self.interpreter.get_tensor(self.output_details[0]['index']),
self.interpreter.get_tensor(self.output_details[1]['index'])
]
out_boxes, out_classes, out_scores = yolo3_postprocess_np(
output_data,
image_shape,
self.anchors,
len(self.class_names),
self.model_image_size,
max_boxes=20,
confidence=0.35)
self.log.info('Found {} boxes for {}'.format(len(out_boxes), 'img'))
end = time.time()
self.log.info("Inference time: {:.8f}s".format(end - start))
if out_classes is None or len(out_classes) == 0:
self.log.warning("No boxes found!")
return image_data, None, None
order = out_boxes[:, 0].argsort()
out_boxes = out_boxes[order]
out_classes = out_classes[order]
out_scores = out_scores[order]
yref_top = min(out_boxes[:, 1])
yref_top_idx = np.argmin(out_boxes[:, 1])
r_number, r_screen = ([] for i in range(2))
for idx, box in enumerate(out_boxes):
_, ymin, _, ymax = box
if ymin < 195:
r_screen.append(out_classes[idx])
else:
r_number.append(out_classes[idx])
#r_number = int("".join("{0}".format(n) for n in r_number))
r_number = "".join(str(n) for n in r_number)
r_screen = "".join(str(n) for n in r_screen)
#draw result on input image
image_array = np.array(image, dtype='uint8')
image_array = draw_boxes(image_array, out_boxes, out_classes,
out_scores, self.class_names, self.colors)
return Image.fromarray(image_array), r_number, r_screen
def get_prediction_class_records(model, model_format, annotation_records, anchors, class_names, model_image_size, conf_threshold, save_result):
'''
Do the predict with YOLO model on annotation images to get predict class dict
predict class dict would contain image_name, coordinary and score, and
sorted by score:
pred_classes_records = {
'car': [
['000001.jpg','94,115,203,232',0.98],
['000002.jpg','82,64,154,128',0.93],
...
],
...
}
'''
if model_format == 'MNN':
#MNN inference engine need create session
session = model.createSession()
# create txt file to save prediction result, with
# save format as annotation file but adding score, like:
#
# path/to/img1.jpg 50,100,150,200,0,0.86 30,50,200,120,3,0.95
#
os.makedirs('result', exist_ok=True)
result_file = open(os.path.join('result','detection_result.txt'), 'w')
pred_classes_records = OrderedDict()
pbar = tqdm(total=len(annotation_records), desc='Eval model')
for (image_name, gt_records) in annotation_records.items():
image = Image.open(image_name)
if image.mode != 'RGB':
image = image.convert('RGB')
image_array = np.array(image, dtype='uint8')
# support of tflite model
if model_format == 'TFLITE':
pred_boxes, pred_classes, pred_scores = yolo_predict_tflite(model, image, anchors, len(class_names), conf_threshold)
# support of MNN model
elif model_format == 'MNN':
pred_boxes, pred_classes, pred_scores = yolo_predict_mnn(model, session, image, anchors, len(class_names), conf_threshold)
# support of TF 1.x frozen pb model
elif model_format == 'PB':
pred_boxes, pred_classes, pred_scores = yolo_predict_pb(model, image, anchors, len(class_names), model_image_size, conf_threshold)
# support of ONNX model
elif model_format == 'ONNX':
pred_boxes, pred_classes, pred_scores = yolo_predict_onnx(model, image, anchors, len(class_names), conf_threshold)
# normal keras h5 model
elif model_format == 'H5':
pred_boxes, pred_classes, pred_scores = yolo_predict_keras(model, image, anchors, len(class_names), model_image_size, conf_threshold)
else:
raise ValueError('invalid model format')
#print('Found {} boxes for {}'.format(len(pred_boxes), image_name))
pbar.update(1)
# save prediction result to txt
result_file.write(image_name)
for box, cls, score in zip(pred_boxes, pred_classes, pred_scores):
xmin, ymin, xmax, ymax = box
box_annotation = " %d,%d,%d,%d,%d,%f" % (
xmin, ymin, xmax, ymax, cls, score)
result_file.write(box_annotation)
result_file.write('\n')
result_file.flush()
if save_result:
gt_boxes, gt_classes, gt_scores = transform_gt_record(gt_records, class_names)
result_dir=os.path.join('result','detection')
os.makedirs(result_dir, exist_ok=True)
colors = get_colors(class_names)
image_array = draw_boxes(image_array, gt_boxes, gt_classes, gt_scores, class_names, colors=None, show_score=False)
image_array = draw_boxes(image_array, pred_boxes, pred_classes, pred_scores, class_names, colors)
image = Image.fromarray(image_array)
# here we handle the RGBA image
if(len(image.split()) == 4):
r, g, b, a = image.split()
image = Image.merge("RGB", (r, g, b))
image.save(os.path.join(result_dir, image_name.split(os.path.sep)[-1]))
# Nothing detected
if pred_boxes is None or len(pred_boxes) == 0:
continue
for box, cls, score in zip(pred_boxes, pred_classes, pred_scores):
pred_class_name = class_names[cls]
xmin, ymin, xmax, ymax = box
coordinate = "{},{},{},{}".format(xmin, ymin, xmax, ymax)
#append or add predict class item
record = [os.path.basename(image_name), coordinate, score]
if pred_class_name in pred_classes_records:
pred_classes_records[pred_class_name].append(record)
else:
pred_classes_records[pred_class_name] = list([record])
# sort pred_classes_records for each class according to score
for pred_class_list in pred_classes_records.values():
pred_class_list.sort(key=lambda ele: ele[2], reverse=True)
pbar.close()
result_file.close()
return pred_classes_records
def get_prediction_class_records(model, model_format, annotation_records,
anchors, class_names, model_image_size,
conf_threshold, save_result):
'''
Do the predict with YOLO model on annotation images to get predict class dict
predict class dict would contain image_name, coordinary and score, and
sorted by score:
pred_classes_records = {
'car': [
['00001.jpg','94,115,203,232',0.98],
['00002.jpg','82,64,154,128',0.93],
...
],
...
}
'''
if model_format == 'MNN':
#MNN inference engine need create session
session = model.createSession()
pred_classes_records = {}
pbar = tqdm(total=len(annotation_records), desc='Eval model')
for (image_name, gt_records) in annotation_records.items():
image = Image.open(image_name)
image_array = np.array(image, dtype='uint8')
# support of tflite model
if model_format == 'TFLITE':
pred_boxes, pred_classes, pred_scores = yolo_predict_tflite(
model, image, anchors, len(class_names), conf_threshold)
# support of MNN model
elif model_format == 'MNN':
pred_boxes, pred_classes, pred_scores = yolo_predict_mnn(
model, session, image, anchors, len(class_names),
conf_threshold)
# support of TF 1.x frozen pb model
elif model_format == 'PB':
pred_boxes, pred_classes, pred_scores = yolo_predict_pb(
model, image, anchors, len(class_names), model_image_size,
conf_threshold)
# normal keras h5 model
elif model_format == 'H5':
pred_boxes, pred_classes, pred_scores = yolo_predict_keras(
model, image, anchors, len(class_names), model_image_size,
conf_threshold)
else:
raise ValueError('invalid model format')
#print('Found {} boxes for {}'.format(len(pred_boxes), image_name))
pbar.update(1)
if save_result:
gt_boxes, gt_classes, gt_scores = transform_gt_record(
gt_records, class_names)
result_dir = os.path.join('result', 'detection')
touchdir(result_dir)
colors = get_colors(class_names)
image_array = draw_boxes(image_array,
gt_boxes,
gt_classes,
gt_scores,
class_names,
colors=None,
show_score=False)
image_array = draw_boxes(image_array, pred_boxes, pred_classes,
pred_scores, class_names, colors)
image = Image.fromarray(image_array)
# here we handle the RGBA image
if (len(image.split()) == 4):
r, g, b, a = image.split()
image = Image.merge("RGB", (r, g, b))
image.save(
os.path.join(result_dir,
image_name.split(os.path.sep)[-1]))
# Nothing detected
if pred_boxes is None or len(pred_boxes) == 0:
continue
for box, cls, score in zip(pred_boxes, pred_classes, pred_scores):
pred_class_name = class_names[cls]
xmin, ymin, xmax, ymax = box
coordinate = "{},{},{},{}".format(xmin, ymin, xmax, ymax)
#append or add predict class item
if pred_class_name in pred_classes_records:
pred_classes_records[pred_class_name].append(
[image_name, coordinate, score])
else:
pred_classes_records[pred_class_name] = list(
[[image_name, coordinate, score]])
# sort pred_classes_records for each class according to score
for pred_class_list in pred_classes_records.values():
pred_class_list.sort(key=lambda ele: ele[2], reverse=True)
pbar.close()
return pred_classes_records
def detect_image(self, image, apply_constraints=False):
if self.model_image_size != (None, None):
assert self.model_image_size[
0] % 32 == 0, 'Multiples of 32 required'
assert self.model_image_size[
1] % 32 == 0, 'Multiples of 32 required'
image_data = preprocess_image(image, self.model_image_size)
# prepare origin image shape, (height, width) format
image_shape = np.array([image.size[1], image.size[0]])
image_shape = np.expand_dims(image_shape, 0)
start = time.time()
out_boxes, out_classes, out_scores = self.predict(
image_data, image_shape)
end = time.time()
print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
print("Inference time: {:.8f}s".format(end - start))
if (apply_constraints):
print('Applying constraints...')
idx_stem, = np.where(self.class_names == 'stem')
idx_stick, = np.where(self.class_names == 'stick')
# Stem constraint
if (len(idx_stick) > 0):
stems, = np.where(out_classes == idx_stem)
if (len(stems) > 1):
higher_weighted_dist = -9999
higher_weighted_dist_idx = -1
for i in stems:
# Calculating the weighted difference between the stick and the current stem
bottom_dist = abs(out_boxes[i][-1] -
out_boxes[idx_stick][-1])
dist = min([
abs(out_boxes[idx_stick][0] - out_boxes[i][1]),
abs(out_boxes[idx_stick][1] - out_boxes[i][0])
])
weighted_dist = bottom_dist * (
1 / (dist + 0.00001)) + bottom_dist * out_scores[i]
if (weighted_dist > higher_weighted_dist):
higher_weighted_dist = weighted_dist
if (higher_weighted_dist_idx >= 0):
out_boxes = np.delete(out_boxes,
higher_weighted_dist_idx,
axis=0)
out_classes = np.delete(
out_classes,
higher_weighted_dist_idx,
axis=0)
out_scores = np.delete(
out_scores,
higher_weighted_dist_idx,
axis=0)
higher_weighted_dist_idx = i
#draw result on input image
image_array = np.array(image, dtype='uint8')
image_array = draw_boxes(image_array, out_boxes, out_classes,
out_scores, self.class_names, self.colors)
return Image.fromarray(image_array)
elapsed = time.time() - start
fps = 1 / elapsed
print('Inference time in ms: %.2f' % (elapsed * 1000))
# post-processing
image_shape = tuple((frame.shape[0], frame.shape[1]))
out_boxes, out_classes, out_scores = yolo3_postprocess_np(
detected_boxes,
image_shape,
anchors,
len(labels), (input_height, input_width),
max_boxes=10,
confidence=conf_threshold,
iou_threshold=iou_threshold,
elim_grid_sense=True)
# modified draw_boxes function to return an openCV formatted image
img_bbox = draw_boxes(img, out_boxes, out_classes, out_scores, labels,
colors)
# draw information overlay onto the frames
cv2.putText(img_bbox,
'Inference Running on : {0}'.format(backend_name),
(30, 50), font, font_size, color, font_thickness)
cv2.putText(
img_bbox, 'FPS : {0} | Inference Time : {1}ms'.format(
int(fps), round((elapsed * 1000), 2)), (30, 80), font,
font_size, color, font_thickness)
if input_mode in 'directory':
out_file = "detections_{0}.jpg".format(image_id)
out_file = os.path.join(result_dir, out_file)
cv2.imwrite(out_file, img_bbox)
if input_mode in 'image':
cv2.imwrite("detections.jpg", img_bbox)
print("Output image is saved in detections.jpg")
