frcnn_evaluate.py 7.57 KB
Newer Older
sjjsmuel's avatar
sjjsmuel committed
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192
import tensorflow as tf
from pathlib import Path
import numpy as np

from PIL import Image
import cv2
from helpers.AnnotationLocationLoader import AnnotationLocationLoader
from helpers.PredictionLocationLoader import PredictionLocationLoader
from optparse import OptionParser


def run_inference_for_single_image(model, image):
    '''
    Inference function for single image based on function from TF Object Detection API
    https://github.com/tensorflow/models/blob/master/research/object_detection/colab_tutorials/object_detection_tutorial.ipynb
    '''
    image = np.asarray(image)
    # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
    input_tensor = tf.convert_to_tensor(image)
    # The model expects a batch of images, so add an axis with `tf.newaxis`.
    input_tensor = input_tensor[tf.newaxis, ...]

    # Run inference
    model_fn = model.signatures['serving_default']
    output_dict = model_fn(input_tensor)

    # All outputs are batches tensors.
    # Convert to numpy arrays, and take index [0] to remove the batch dimension.
    # We're only interested in the first num_detections.
    num_detections = int(output_dict.pop('num_detections'))
    output_dict = {key: value[0, :num_detections].numpy()
                   for key, value in output_dict.items()}
    output_dict['num_detections'] = num_detections

    # detection_classes should be ints.
    output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64)

    return output_dict


def save_inference(model, image_path, output_path, annotation_loader):
    """
    Run the inference for an image, predict localizations, create visualisations and save everything to file
    :param model: Faster R-CNN model
    :param image_path: path the image for which to run the inference
    :param output_path: path to base output folder
    :param annotation_loader: helper to get the corresponding ground truth
    :return: None - save results to disk
    """
    # get image as numpy array
    filename = image_path.name
    image_np = np.array(Image.open(image_path))
    # run inference
    output_dict = run_inference_for_single_image(model, image_np)

    selected_indices = tf.image.non_max_suppression(output_dict['detection_boxes'], output_dict['detection_scores'], 20, 0.3)

    img = cv2.imread(str(image_path), cv2.IMREAD_IGNORE_ORIENTATION  | cv2.IMREAD_COLOR)
    height, width, _ = img.shape

    # draw the annotated labels to the original picture
    annotations = annotation_loader.get_annotations(filename, ['caries'])
    for annotation in annotations:
        cv2.rectangle(img, annotation[1][0], annotation[1][1], (0, 255, 0), 3)

    prepared_boxes = []
    for i, detection_score in enumerate(output_dict['detection_scores']):
        if i not in selected_indices:
            continue
        if detection_score < 0.5:
            continue
        # ymin, xmin, ymax, xmax as relative to the image.
        bbox = output_dict['detection_boxes'][i]
        ymin = int(round(bbox[0] * height))
        xmin = int(round(bbox[1] * width))
        ymax = int(round(bbox[2] * height))
        xmax = int(round(bbox[3] * width))

        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (255, 255, 255), 3)
        class_prediction = output_dict['detection_classes'][i]
        class_prediction = index_class_map[class_prediction]
        class_score = output_dict['detection_scores'][i]
        label = "{} {:.2f}%".format(class_prediction, (class_score*100))
        cv2.putText(img, label, (xmin, (ymin - 5)), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

        box = (class_prediction, [(xmin, ymin), (xmax, ymax), class_score])
        prepared_boxes.append(box)

    with open(str(output_path.parent/"predictions.txt"), "a") as predictions_file:
        predictions_file.write('{}; {}\n'.format(filename, str(prepared_boxes)))

    cv2.imwrite(str(output_path/filename), img)


parser = OptionParser()

parser.add_option("-i", "--path_evaluation", dest="evaluation_path", help="Path to evaluation input.", default="./input/evaluation_data")
parser.add_option("--model_folder", dest="model_folder", help="Path to the model to evaluate.", default="./out/trained_models/frcnn_model/exported_model")
parser.add_option("-o", "--path_output", dest="output_path", help="Path to base folder for output data.", default='./out')

(options, args) = parser.parse_args()

evaluation_folder = Path(options.evaluation_path)
frcnn_checkpoint_folder = Path(options.model_folder)
frcnn_output_folder = Path(options.output_path) / 'evaluation_frcnn'
if not frcnn_output_folder.exists():
    frcnn_output_folder.mkdir()

class_index_map = {'caries': 1, 'no_caries': 2}
index_class_map = {}
for element in class_index_map:
    index_class_map[class_index_map[element]] = element

annotation_loader = AnnotationLocationLoader(images_base_folder=evaluation_folder)

# clean up predicted boxes from last run
with open(str(frcnn_output_folder/"predictions.txt"), "w") as predictions_file:
    predictions_file.write("")

model_dir = frcnn_checkpoint_folder / "saved_model"
model = tf.saved_model.load(str(model_dir))

# run inference for all images in the evaluation folder
for path in [path for path in evaluation_folder.iterdir() if path.is_dir()]:
    # define an output folder for each class
    out_path = Path(frcnn_output_folder) / path.name
    if not out_path.exists():
        out_path.mkdir()

    # get the filenames of all test-images in the current folder
    filenames = [item for item in path.glob('*') if item.name != '.DS_Store']

    for img in filenames:
        save_inference(model, img, out_path, annotation_loader)
        print("[INFO] Finished {}/{}".format(img.parent.name, img.name))

# reload currently saved predictions for evaluation the classification
prediction_loader = PredictionLocationLoader(prediction_file=str(frcnn_output_folder/'predictions.txt'), images_base_folder=evaluation_folder)

tp = 0
fp = 0
fn = 0
tn = 0
fp_list = []
fn_list = []
# evaluate the classification capabilities of the Faster R-CNN model
for path in [path for path in evaluation_folder.iterdir() if path.is_dir()]:
    for image in [item for item in path.glob('*') if item.name != '.DS_Store']:
        predictions = prediction_loader.get_annotations(str(image.name))
        if len(predictions) > 0: #prediction is positive
            if path.name == 'caries':
                tp += 1
            else:
                fp += 1
                fp_list.append(str(image.name))
        else:   # prediction is negative
            if path.name == 'caries':
                fn += 1
            else:
                tn += 1
                fn_list.append(str(image.name))

p = float(tp) / (tp + fp)
r = float(tp) / (tp + fn)
f1 = (2 * p * r) / (p + r)

# save and print the classification results
with open(str(frcnn_output_folder/'classification_results.txt'), 'w') as out_file:
    out_file.write("Klassifikationsergebnisse für " + str(frcnn_checkpoint_folder)+"\n")
    out_file.write("\n")
    out_file.write("tp " + str(tp)+"\n")
    out_file.write("fp " + str(fp)+"\n")
    out_file.write("fn " + str(fn)+"\n")
    out_file.write("tn " + str(tn)+"\n")
    out_file.write("Precision " + str(p)+"\n")
    out_file.write("Recall " + str(r)+"\n")
    out_file.write("F1-Measure " + str(f1)+"\n")
    out_file.write("\n")
    out_file.write("False Positive Bilder:\n")
    out_file.write(str(fp_list))
    out_file.write("\n")
    out_file.write("False Negative Bilder:\n")
    out_file.write(str(fn_list))

print('Results for the Faster R-CNN as Classifier:')
print("Precision " + str(p))
print("Recall " + str(r))
print("F1-Measure " + str(f1))
print(tp, fn)
print(fp, tn)

print("[INFO] Evaluation finished.")