File size: 4,833 Bytes
26364eb |
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 |
from peekingduck.pipeline.nodes.model import yolo as pkd_yolo
import cv2
from collections import defaultdict
import numpy as np
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
def convert_labels(inference_labels_dict, bbox_labels):
for k, v in inference_labels_dict.items():
bbox_labels[bbox_labels == k] = v
# FutureWarning: elementwise comparison failed; returning scalar, but in the future will perform elementwise comparison
# throws up this warning because making a change string to int is something that numpy disagrees with (???).
return bbox_labels
def process_masks(inference_outputs, inference_labels_dict):
mask_labels = convert_labels(inference_labels_dict, inference_outputs["bbox_labels"])
masks = inference_outputs["masks"]
scores = inference_outputs['bbox_scores']
processed_output = [[masks[i], scores[i], int(mask_labels[i])] for i in range(len(scores))]
return processed_output
def process_bboxes(inference_outputs, inference_labels_dict):
bbox_labels = inference_outputs["bbox_labels"]
bbox_labels = convert_labels(inference_labels_dict, bbox_labels)
bboxes = inference_outputs["bboxes"]
bbox_scores = inference_outputs["bbox_scores"]
# stack the bbox_scores and bbox_labels
# hence, array(['score', 'score','score']) and array(['class','class','class'])
# becomes array([['score','class'], ['score','class'],['score','class']])
stacked = np.stack((bbox_scores, bbox_labels), axis = 1)
# concatenate the values of the bbox wih the stacked values above
# use concatenate here because it is 1xnxm with 1xnxl dimension so it works
# it's just maths, people!
concated = np.concatenate((bboxes, stacked), axis = 1)
return concated.astype(np.float32)
def run_inference(img_matrix, model, inference_labels_dict = {'person': 1, 'bicycle': 2}, task = "det"):
"""Helper function to run per image inference, get bbox, labels and scores and stack them for confusion matrix output
Args:
img_matrix (np.array): _description_
model: _description_
labels_dict (dict, optional): _description_. Defaults to {'person': 0, 'bicycle': 1}.
Returns:
concated (np.array): concatenated inference of n x (bbox (default is x1, y1, x2, y2), score, class)
img_matrix.shape (np vector): vector with [Height * Weight * Dimension] values
"""
# print(img_matrix.shape)
# for img_matrix, it's HxWxD. Need to resize it for the confusion matrix
# modify this to change the run to your model's inference method eg model(img) in pytorch
inference_inputs = {"img": img_matrix}
inference_outputs = model.run(inference_inputs)
# pkd outputs for segmentation - {'bboxes': [[],[]..,[]], 'bbox_labels':[], 'bbox_scores':[], 'masks':[[[],[],[]]]}
if task == "seg":
processed_output = process_masks(inference_outputs, inference_labels_dict)
elif task == "det":
processed_output = process_bboxes(inference_outputs, inference_labels_dict)
return processed_output, img_matrix.shape
class Inference:
def __init__(self, model, cfg_obj):
self.model = model
self.labels_dict = cfg_obj['error_analysis']['labels_dict']
self.inference_labels_dict = cfg_obj['error_analysis']['inference_labels_dict']
self.task = cfg_obj['error_analysis']['task']
def run_inference_path(self, img_path):
"""use if img_path is specified
Args:
img_path (_type_): _description_
Returns:
_type_: _description_
"""
image_orig = cv2.imread(img_path)
image_orig = cv2.cvtColor(image_orig, cv2.COLOR_BGR2RGB)
output = run_inference(image_orig, self.model, inference_labels_dict = self.inference_labels_dict, task = self.task)
return output
def run_inference_byte(self, img_bytes):
"""use if the img_bytes is passed in instead of path
Args:
img_bytes (_type_): _description_
Returns:
_type_: _description_
"""
img_decoded = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), -1)
img_decoded = cv2.cvtColor(img_decoded, cv2.COLOR_BGR2RGB)
output = run_inference(img_decoded, self.model, labels_dict = self.inference_labels_dict, task = self.task)
return output
if __name__ == "__main__":
import yaml
from src.error_analysis import load_model
cfg_file = open("cfg/cfg.yml")
cfg_obj = yaml.load(cfg_file, Loader=yaml.FullLoader)
img_path = "./data/annotations_trainval2017/coco_person/000000576052.jpg"
inference_obj = Inference(model = load_model(cfg_obj), cfg_obj = cfg_obj)
output = inference_obj.run_inference_path(img_path)
print (output) |