all_detections = [[None for i in range(generator.num_classes()) if generator.has_label(i)] for j in range(generator.size())]
for i in range(generator.size()):
raw_image = generator.load_image(i)
image = generator.preprocess_image(raw_image.copy())
image, scale = generator.resize_image(image)
if keras.backend.image_data_format() == "channels_first":
image = image.transpose((2, 0, 1))
// run network
boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))[:3]
// correct boxes for image scale
boxes /= scale
// select indices which have a score above the threshold
indices = np.where(scores[0, :] > score_threshold)[0]
// select those scores
scores = scores[0][indices]
// find the order with which to sort the scores
scores_sort = np.argsort(-scores)[:max_detections]
// select detections
image_boxes = boxes[0, indices[scores_sort], :]
image_scores = scores[scores_sort]
image_labels = labels[0, indices[scores_sort]]
image_detections = np.concatenate([image_boxes, np.expand_dims(image_scores, axis=1), np.expand_dims(image_labels, axis=1)], axis=1)
if save_path is not None:
draw_annotations(raw_image, generator.load_annotations(i), label_to_name=generator.label_to_name)
draw_detections(raw_image, image_boxes, image_scores, image_labels, label_to_name=generator.label_to_name)
cv2.imwrite(os.path.join(save_path, "{}.png".format(i)), raw_image)
// copy detections to all_detections
for label in range(generator.num_classes()):
if not generator.has_label(label):
continue
all_detections[i][label] = image_detections[image_detections[:, -1] == label, :-1]
print("{}/{}".format(i + 1, generator.size()), end="\r")
return all_detections
After Change
all_detections = [[None for i in range(generator.num_classes()) if generator.has_label(i)] for j in range(generator.size())]
for i in progressbar.progressbar(range(generator.size()), prefix="Running network: "):
raw_image = generator.load_image(i)
image = generator.preprocess_image(raw_image.copy())
image, scale = generator.resize_image(image)
