Spaces:
Sleeping
Sleeping
| 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(labels_dict, bbox_labels): | |
| for k, v in 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 run_inference(img_matrix, model, labels_dict = {'person': 1, 'bicycle': 2}): | |
| """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 | |
| inference_inputs = {"img": img_matrix} | |
| # modify this to change the run to your model's inference method eg model(img) in pytorch | |
| inference_outputs = model.run(inference_inputs) | |
| bbox_labels = inference_outputs["bbox_labels"] | |
| bbox_labels = convert_labels(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), img_matrix.shape | |
| class Inference: | |
| def __init__(self, model, cfg_obj): | |
| self.model = model | |
| self.labels_dict = cfg_obj['error_analysis']['labels_dict'] | |
| 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, labels_dict = self.labels_dict) | |
| 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.labels_dict) | |
| return output | |
| if __name__ == "__main__": | |
| import yaml | |
| cfg_file = open(cfg_path) | |
| cfg_obj = yaml.load(cfg_file, Loader=yaml.FullLoader) | |
| img_path = "./data/annotations_trainval2017/coco_person/000000000139.jpg" | |
| inference_obj = Inference(model = pkd_yolo.Node(model_type = "v4tiny", detect= ["Person"] , cfg_obj = cfg_obj)) | |
| print (inference_obj.run_inference_path(img_path)) |