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)