NPU人体检测

py_utils/coco_utils.py

@@ -0,0 +1,176 @@
+from copy import copy
+import os
+import cv2
+import numpy as np
+import json
+
+class Letter_Box_Info():
+    def __init__(self, shape, new_shape, w_ratio, h_ratio, dw, dh, pad_color) -> None:
+        self.origin_shape = shape
+        self.new_shape = new_shape
+        self.w_ratio = w_ratio
+        self.h_ratio = h_ratio
+        self.dw = dw 
+        self.dh = dh
+        self.pad_color = pad_color
+
+
+def coco_eval_with_json(anno_json, pred_json):
+    from pycocotools.coco import COCO
+    from pycocotools.cocoeval import COCOeval
+    anno = COCO(anno_json)
+    pred = anno.loadRes(pred_json)
+    eval = COCOeval(anno, pred, 'bbox')
+    # eval.params.useCats = 0
+    # eval.params.maxDets = list((100, 300, 1000))
+    # a = np.array(list(range(50, 96, 1)))/100
+    # eval.params.iouThrs = a
+    eval.evaluate()
+    eval.accumulate()
+    eval.summarize()
+    map, map50 = eval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)
+
+    print('map  --> ', map)
+    print('map50--> ', map50)
+    print('map75--> ', eval.stats[2])
+    print('map85--> ', eval.stats[-2])
+    print('map95--> ', eval.stats[-1])
+
+class COCO_test_helper():
+    def __init__(self, enable_letter_box = False) -> None:
+        self.record_list = []
+        self.enable_ltter_box = enable_letter_box
+        if self.enable_ltter_box is True:
+            self.letter_box_info_list = []
+        else:
+            self.letter_box_info_list = None
+
+    def letter_box(self, im, new_shape, pad_color=(0,0,0), info_need=False):
+        # Resize and pad image while meeting stride-multiple constraints
+        shape = im.shape[:2]  # current shape [height, width]
+        if isinstance(new_shape, int):
+            new_shape = (new_shape, new_shape)
+
+        # Scale ratio
+        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+
+        # Compute padding
+        ratio = r  # width, height ratios
+        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
+        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
+
+        dw /= 2  # divide padding into 2 sides
+        dh /= 2
+
+        if shape[::-1] != new_unpad:  # resize
+            im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
+        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+        im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=pad_color)  # add border
+        
+        if self.enable_ltter_box is True:
+            self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, ratio, ratio, dw, dh, pad_color))
+        if info_need is True:
+            return im, ratio, (dw, dh)
+        else:
+            return im
+
+    def direct_resize(self, im, new_shape, info_need=False):
+        shape = im.shape[:2]
+        h_ratio = new_shape[0]/ shape[0]
+        w_ratio = new_shape[1]/ shape[1]
+        if self.enable_ltter_box is True:
+            self.letter_box_info_list.append(Letter_Box_Info(shape, new_shape, w_ratio, h_ratio, 0, 0, (0,0,0)))
+        im = cv2.resize(im, (new_shape[1], new_shape[0]))
+        return im
+
+    def get_real_box(self, box, in_format='xyxy'):
+        bbox = copy(box)
+        if self.enable_ltter_box == True:
+        # unletter_box result
+            if in_format=='xyxy':
+                bbox[:,0] -= self.letter_box_info_list[-1].dw
+                bbox[:,0] /= self.letter_box_info_list[-1].w_ratio
+                bbox[:,0] = np.clip(bbox[:,0], 0, self.letter_box_info_list[-1].origin_shape[1])
+
+                bbox[:,1] -= self.letter_box_info_list[-1].dh
+                bbox[:,1] /= self.letter_box_info_list[-1].h_ratio
+                bbox[:,1] = np.clip(bbox[:,1], 0, self.letter_box_info_list[-1].origin_shape[0])
+
+                bbox[:,2] -= self.letter_box_info_list[-1].dw
+                bbox[:,2] /= self.letter_box_info_list[-1].w_ratio
+                bbox[:,2] = np.clip(bbox[:,2], 0, self.letter_box_info_list[-1].origin_shape[1])
+
+                bbox[:,3] -= self.letter_box_info_list[-1].dh
+                bbox[:,3] /= self.letter_box_info_list[-1].h_ratio
+                bbox[:,3] = np.clip(bbox[:,3], 0, self.letter_box_info_list[-1].origin_shape[0])
+        return bbox
+
+    def get_real_seg(self, seg):
+        #! fix side effect
+        dh = int(self.letter_box_info_list[-1].dh)
+        dw = int(self.letter_box_info_list[-1].dw)
+        origin_shape = self.letter_box_info_list[-1].origin_shape
+        new_shape = self.letter_box_info_list[-1].new_shape
+        if (dh == 0) and (dw == 0) and origin_shape == new_shape:
+            return seg
+        elif dh == 0 and dw != 0:
+            seg = seg[:, :, dw:-dw] # a[0:-0] = []
+        elif dw == 0 and dh != 0 : 
+            seg = seg[:, dh:-dh, :]
+        seg = np.where(seg, 1, 0).astype(np.uint8).transpose(1,2,0)
+        seg = cv2.resize(seg, (origin_shape[1], origin_shape[0]), interpolation=cv2.INTER_LINEAR)
+        if len(seg.shape) < 3:
+            return seg[None,:,:]
+        else:
+            return seg.transpose(2,0,1)
+
+    def add_single_record(self, image_id, category_id, bbox, score, in_format='xyxy', pred_masks = None):
+        if self.enable_ltter_box == True:
+        # unletter_box result
+            if in_format=='xyxy':
+                bbox[0] -= self.letter_box_info_list[-1].dw
+                bbox[0] /= self.letter_box_info_list[-1].w_ratio
+
+                bbox[1] -= self.letter_box_info_list[-1].dh
+                bbox[1] /= self.letter_box_info_list[-1].h_ratio
+
+                bbox[2] -= self.letter_box_info_list[-1].dw
+                bbox[2] /= self.letter_box_info_list[-1].w_ratio
+
+                bbox[3] -= self.letter_box_info_list[-1].dh
+                bbox[3] /= self.letter_box_info_list[-1].h_ratio
+                # bbox = [value/self.letter_box_info_list[-1].ratio for value in bbox]
+
+        if in_format=='xyxy':
+        # change xyxy to xywh
+            bbox[2] = bbox[2] - bbox[0]
+            bbox[3] = bbox[3] - bbox[1]
+        else:
+            assert False, "now only support xyxy format, please add code to support others format"
+        
+        def single_encode(x):
+            from pycocotools.mask import encode
+            rle = encode(np.asarray(x[:, :, None], order="F", dtype="uint8"))[0]
+            rle["counts"] = rle["counts"].decode("utf-8")
+            return rle
+
+        if pred_masks is None:
+            self.record_list.append({"image_id": image_id,
+                                    "category_id": category_id,
+                                    "bbox":[round(x, 3) for x in bbox],
+                                    'score': round(score, 5),
+                                    })
+        else:
+            rles = single_encode(pred_masks)
+            self.record_list.append({"image_id": image_id,
+                                    "category_id": category_id,
+                                    "bbox":[round(x, 3) for x in bbox],
+                                    'score': round(score, 5),
+                                    'segmentation': rles,
+                                    })
+    
+    def export_to_json(self, path):
+        with open(path, 'w') as f:
+            json.dump(self.record_list, f)
+

py_utils/onnx_executor.py

@@ -0,0 +1,103 @@
+import os
+import numpy as np
+import onnxruntime as rt
+
+type_map = {
+    'tensor(int32)' : np.int32,
+    'tensor(int64)' : np.int64,
+    'tensor(float32)' : np.float32,
+    'tensor(float64)' : np.float64,
+    'tensor(float)' : np.float32,
+}
+if getattr(np, 'bool', False):
+    type_map['tensor(bool)'] = np.bool
+else:
+    type_map['tensor(bool)'] = bool
+
+def ignore_dim_with_zero(_shape, _shape_target):
+    _shape = list(_shape)
+    _shape_target = list(_shape_target)
+    for i in range(_shape.count(1)):
+        _shape.remove(1)
+    for j in range(_shape_target.count(1)):
+        _shape_target.remove(1)
+    if _shape == _shape_target:
+        return True
+    else:
+        return False
+
+
+class ONNX_model_container_py:
+    def __init__(self, model_path) -> None:
+        # sess_options=
+        sp_options = rt.SessionOptions()
+        sp_options.log_severity_level = 3
+        # [1 for info, 2 for warning, 3 for error, 4 for fatal] 
+        self.sess = rt.InferenceSession(model_path, sess_options=sp_options, providers=['CPUExecutionProvider'])
+        self.model_path = model_path
+
+    def run(self, input_datas):
+        if len(input_datas) < len(self.sess.get_inputs()):
+            assert False,'inputs_datas number not match onnx model{} input'.format(self.model_path)
+        elif len(input_datas) > len(self.sess.get_inputs()):
+            print('WARNING: input datas number large than onnx input node')
+
+        input_dict = {}
+        for i, _input in enumerate(self.sess.get_inputs()):
+            # convert type
+            if _input.type in type_map and \
+                type_map[_input.type] != input_datas[i].dtype:
+                print('WARNING: force data-{} from {} to {}'.format(i, input_datas[i].dtype, type_map[_input.type]))
+                input_datas[i] = input_datas[i].astype(type_map[_input.type])
+            
+            # reshape if need
+            if _input.shape != list(input_datas[i].shape):
+                if ignore_dim_with_zero(input_datas[i].shape,_input.shape):
+                    input_datas[i] = input_datas[i].reshape(_input.shape)
+                    print("WARNING: reshape inputdata-{}: from {} to {}".format(i, input_datas[i].shape, _input.shape))
+                else:
+                    assert False, 'input shape{} not match real data shape{}'.format(_input.shape, input_datas[i].shape)
+            input_dict[_input.name] = input_datas[i]
+
+        output_list = []
+        for i in range(len(self.sess.get_outputs())):
+            output_list.append(self.sess.get_outputs()[i].name)
+
+        #forward model
+        res = self.sess.run(output_list, input_dict)
+        return res
+
+
+class ONNX_model_container_cpp:
+    def __init__(self, model_path) -> None:
+        pass
+
+    def run(self, input_datas):
+        pass
+
+
+def ONNX_model_container(model_path, backend='py'):
+    if backend == 'py':
+        return ONNX_model_container_py(model_path)
+    elif backend == 'cpp':
+        return ONNX_model_container_cpp(model_path)
+
+
+def reset_onnx_shape(onnx_model_path, output_path, input_shapes):
+    if isinstance(input_shapes[0], int):
+        command = "python -m onnxsim {} {} --input-shape {}".format(onnx_model_path, output_path, ','.join([str(v) for v in input_shapes]))
+    else:
+        if len(input_shapes)!= 1:
+            print("RESET ONNX SHAPE with more than one input, try to match input name")
+            sess = rt.InferenceSession(onnx_model_path)
+            input_names = [input.name for input in sess.get_inputs()]
+            command = "python -m onnxsim {} {} --input-shape ".format(onnx_model_path, output_path)
+            for i, input_name in enumerate(input_names):
+                command += "{}:{} ".format(input_name, ','.join([str(v) for v in input_shapes[i]]))
+        else:
+            command = "python -m onnxsim {} {} --input-shape {}".format(onnx_model_path, output_path, ','.join([str(v) for v in input_shapes[0]]))
+    
+    print(command)
+    os.system(command)
+    return output_path
+    

py_utils/pytorch_executor.py

@@ -0,0 +1,52 @@
+import torch
+torch.backends.quantized.engine = 'qnnpack'
+
+def multi_list_unfold(tl):
+    def unfold(_inl, target):
+        if not isinstance(_inl, list) and not isinstance(_inl, tuple):
+            target.append(_inl)
+        else:
+            unfold(_inl)
+
+def flatten_list(in_list):
+    flatten = lambda x: [subitem for item in x for subitem in flatten(item)] if type(x) is list else [x]
+    return flatten(in_list)
+
+class Torch_model_container:
+    def __init__(self, model_path, qnnpack=False) -> None:
+        if qnnpack is True:
+            torch.backends.quantized.engine = 'qnnpack'
+
+        #! Backends must be set before load model.
+        self.pt_model = torch.jit.load(model_path)
+        self.pt_model.eval()
+        holdon = 1
+
+    def run(self, input_datas):
+        assert isinstance(input_datas, list), "input_datas should be a list, like [np.ndarray, np.ndarray]"
+
+        input_datas_torch_type = []
+        for _data in input_datas:
+            input_datas_torch_type.append(torch.tensor(_data))
+
+        for i,val in enumerate(input_datas_torch_type):
+            if val.dtype == torch.float64:
+                input_datas_torch_type[i] = input_datas_torch_type[i].float()
+
+        result = self.pt_model(*input_datas_torch_type)
+
+        if isinstance(result, tuple):
+            result = list(result)
+        if not isinstance(result, list):
+            result = [result]
+        
+        result = flatten_list(result)
+
+        for i in range(len(result)):
+            result[i] = torch.dequantize(result[i])
+
+        for i in range(len(result)):
+            # TODO support quantized_output
+            result[i] = result[i].cpu().detach().numpy()
+
+        return result

py_utils/rknn_executor (copy 1).py

@@ -0,0 +1,31 @@
+from rknn.api import RKNN
+
+
+class RKNN_model_container():
+    def __init__(self, model_path, target=None, device_id=None) -> None:
+        rknn = RKNN()
+
+        # Direct Load RKNN Model
+        rknn.load_rknn(model_path)
+
+        print('--> Init runtime environment')
+        if target==None:
+            ret = rknn.init_runtime()
+        else:
+            ret = rknn.init_runtime(target=target, device_id=device_id)
+        if ret != 0:
+            print('Init runtime environment failed')
+            exit(ret)
+        print('done')
+        
+        self.rknn = rknn 
+
+    def run(self, inputs):
+        if isinstance(inputs, list) or isinstance(inputs, tuple):
+            pass
+        else:
+            inputs = [inputs]
+
+        result = self.rknn.inference(inputs=inputs)
+    
+        return result

py_utils/rknn_executor.py

@@ -0,0 +1,26 @@
+from rknnlite.api import RKNNLite as RKNN
+
+class RKNN_model_container():
+    def __init__(self, model_path, target=None, device_id=None) -> None:
+        rknn = RKNN()
+        rknn.load_rknn(model_path)
+        ret = rknn.init_runtime()
+        self.rknn = rknn
+
+    def run(self, inputs):
+        if self.rknn is None:
+            print("ERROR: rknn has been released")
+            return []
+
+        if isinstance(inputs, list) or isinstance(inputs, tuple):
+            pass
+        else:
+            inputs = [inputs]
+
+        result = self.rknn.inference(inputs=inputs)
+
+        return result
+
+    def release(self):
+        self.rknn.release()
+        self.rknn = None