Add application file

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/dms3/v1.0/dms3_mtl_v1.0.pth filter=lfs diff=lfs merge=lfs -text
+data/dms3/v1.0/dms3_mtl_v1.0.onnx filter=lfs diff=lfs merge=lfs -text
+head_detect/models/HeadDetectorv1.6.onnx filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,179 @@
+import os
+
+import gradio as gr
+
+from inference_mtl import inference_xyl
+from inference_video_mtl import inference_videos
+
+
+SAVE_DIR = "/nfs/volume-236-2/eval_model/model_res"
+# SAVE_DIR = "/Users/didi/Desktop/model_res"
+CURRENT_DIR = ''
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+PLOT_DIR = ""
+TASKS = ["ems", "eye", 'mouth', 'glass', 'mask', 'smoke', 'phone', "eyelid_r", "eyelid_l", 'shift_x', 'shift_y', 'expand']
+TASK_IDX = 0
+
+
+def last_img():
+    global TASK_IDX
+    TASK_IDX -= 1
+    if TASK_IDX < 0:
+        TASK_IDX = len(TASKS) - 1
+    plt_path = f"{PLOT_DIR}/{TASKS[TASK_IDX]}.jpg"
+
+    if not os.path.exists(plt_path):
+        return
+
+    return plt_path
+
+
+def next_img():
+    global TASK_IDX
+    TASK_IDX += 1
+    if TASK_IDX >= len(TASKS):
+        TASK_IDX = 0
+    plt_path = f"{PLOT_DIR}/{TASKS[TASK_IDX]}.jpg"
+
+    if not os.path.exists(plt_path):
+        return
+
+    return plt_path
+
+
+def inference_img(inp, engine):
+    inp = inp[:, :, ::-1]
+
+    if engine == "DMS3":
+        drawn = inference_xyl(inp, vis=False, return_drawn=True)
+        drawn = drawn[:, :, ::-1]
+        return drawn
+
+
+def inference_txt(inp, engine):
+    inp = inp[:, :, ::-1]
+
+    if engine == "DMS3":
+        msg_list = inference_xyl(inp, vis=False, return_drawn=False)[-1]
+        return "\n".join(msg_list)
+
+
+def inference_video(video, engine, device_type_, show_score_, syn_plot_):
+    global CURRENT_DIR
+    CURRENT_DIR = os.path.join(SAVE_DIR, engine)
+    os.makedirs(CURRENT_DIR, exist_ok=True)
+
+    base_name = os.path.splitext(os.path.basename(video))[0][:-8]
+
+    video_info = []
+    if device_type_ == "B100":
+        resize = (1280, 720)
+        video_info.append("720")
+    elif device_type_ == "B200":
+        resize = (1280, 960)
+        video_info.append("960")
+    else:
+        resize = None
+
+    is_syn_plot = syn_plot_ == "是"
+    is_show_score = show_score_ == "是"
+
+    video_info.append("1" if is_syn_plot else "0")
+    video_info.append(base_name + "_res.mp4")
+
+    save_video_name = "_".join(video_info)
+    save_plt_dir = base_name + "_plots"
+    save_csv_name = base_name + "_res.csv"
+
+    save_video_path = f"{CURRENT_DIR}/{save_video_name}"
+    save_plt_path = f"{CURRENT_DIR}/{save_plt_dir}"
+    save_csv_path = f"{CURRENT_DIR}/{save_csv_name}"
+
+    global PLOT_DIR, TASK_IDX
+    PLOT_DIR = save_plt_path
+    TASK_IDX = 0
+
+    if os.path.exists(save_video_path) and os.path.exists(save_plt_path):
+        if not is_show_score:
+            return save_video_path, f"{save_plt_path}/{TASKS[TASK_IDX]}.jpg", None
+        elif os.path.exists(save_csv_path):
+            return save_video_path, f"{save_plt_path}/{TASKS[TASK_IDX]}.jpg", save_csv_path
+
+    inference_videos(
+        video, save_dir=SAVE_DIR, detect_mode='second', frequency=0.2, plot_score=True, save_score=is_show_score,
+        syn_plot=is_syn_plot, save_vdo=True, save_img=False, continuous=False, show_res=False, resize=resize,
+        time_delta=1, save_vdo_path=save_video_path, save_plt_path=save_plt_path, save_csv_path=save_csv_path)
+
+    if not is_show_score:
+        return save_video_path, f"{save_plt_path}/{TASKS[TASK_IDX]}.jpg", None
+
+    return save_video_path, f"{save_plt_path}/{TASKS[TASK_IDX]}.jpg", save_csv_path
+
+
+def reset_state():
+    return None, None, None, None
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""<h1 align="center">算法可视化Demo</h1>""")
+
+    with gr.Tab("Text"):  # 标签页1
+        with gr.Row():    # 并行显示，可开多列
+            with gr.Column():  # 并列显示，可开多行
+                input_img1 = gr.Image(label="", show_label=False)
+                engine1 = gr.Dropdown(["DMS3", "EMS", "TURN2"], label="算法引擎", value="DMS3")
+                btn_txt = gr.Button(value="Submit", label="Submit Image", variant="primary")
+                btn_clear1 = gr.Button("Clear")
+
+            out1 = gr.Text(label="")
+
+    btn_txt.click(inference_txt, inputs=[input_img1, engine1], outputs=out1, show_progress=True)  # 触发
+    btn_clear1.click(reset_state, inputs=[], outputs=[input_img1, out1])
+
+    with gr.Tab("Image"):  # 标签页1
+        with gr.Row():  # 并行显示，可开多列
+            with gr.Column():  # 并列显示，可开多行
+                input_img2 = gr.Image(label="", show_label=False)
+                engine2 = gr.Dropdown(["DMS3", "EMS", "TURN2"], label="算法引擎", value="DMS3")
+                btn_img = gr.Button(value="Submit", label="Submit Image", variant="primary")
+                btn_clear2 = gr.Button("Clear")
+
+            out2 = gr.Image(label="", show_label=False)
+
+    btn_img.click(inference_img, inputs=[input_img2, engine2], outputs=out2, show_progress=True)  # 触发
+    btn_clear2.click(reset_state, inputs=[], outputs=[input_img2, out2])
+
+    with gr.Tab("Video"):  # 标签页2
+        with gr.Row():  # 并行显示，可开多列
+
+            with gr.Column():  # 并列显示，可开多行
+                input_vdo = gr.Video(label="", show_label=False)
+                engine3 = gr.Dropdown(["DMS3", "EMS", "TURN2"], label="算法引擎", value="DMS3")
+                device_type = gr.Radio(["原始", "B100", "B200"], value="原始", label="Device Type")  # 单选
+
+                with gr.Row():
+                    show_score = gr.Radio(["是", "否"], value="是", label="分数明细")  # 单选
+                    syn_plot = gr.Radio(["是", "否"], value="是", label="Syn Plot")  # 单选
+
+                btn_vdo = gr.Button(value="Submit", label="Submit Video", variant="primary")
+                btn_clear3 = gr.Button("Clear")
+
+            with gr.Column():
+                out3 = gr.PlayableVideo(label="", show_label=False)
+                out3_plt = gr.Image(label=f"{TASKS[TASK_IDX]}", show_label=False)
+
+                with gr.Row():
+                    btn_before = gr.Button(value="上一张", label="before")
+                    btn_next = gr.Button(value="下一张", label="next", variant="primary")
+
+        out3_df = gr.DataFrame()
+
+    btn_vdo.click(inference_video, inputs=[input_vdo, engine3, device_type, show_score, syn_plot],
+                  outputs=[out3, out3_plt, out3_df], show_progress=True)  # 触发
+    btn_clear3.click(reset_state, inputs=[], outputs=[input_vdo, out3, out3_plt, out3_df])
+
+    btn_before.click(last_img, inputs=[], outputs=out3_plt)
+    btn_next.click(next_img, inputs=[], outputs=out3_plt)
+
+demo.queue().launch(share=True, inbrowser=True, favicon_path="files/dms3_icon.png")

data/dms3/v1.0/dms3_mtl_v1.0.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4f06fd029196afea2522b194b1ba292a10a31f2792ca5c3d12f35ffd5530f02
+size 2176811

data/dms3/v1.0/dms3_mtl_v1.0.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:483c1b16abc39462272a35006cc92993d98c385f8d05cc2ad869ad5d041b17e7
+size 2444005

head_detect/demo.py

@@ -0,0 +1,59 @@
+# -*- coding:utf-8 –*-
+import os
+import sys
+
+import cv2
+
+sys.path.append(os.path.dirname(os.path.dirname(sys.path[0])))
+from head_detect.head_detector.head_detectorv4 import HeadDetector
+from head_detect.utils_quailty_assurance.utils_quailty_assurance import find_files, write_json
+from utils.images import crop_face_square_rate
+import numpy as np
+
+ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
+model_path = f"{ROOT_DIR}/models/HeadDetectorv1.6.onnx"
+
+headDetector = HeadDetector(onnx_path=model_path)
+
+
+def detect_driver_face(img, show_res=False, show_crop=False):
+    if isinstance(img, str):
+        img = cv2.imread(img)
+    else:
+        img = img.copy()
+    image_heigth,image_width = img.shape[:2]
+    # origin_image = cv2.resize(img, (1280, 720))
+    # origin_image = img[:, w//2:, :]
+    # short_side = min(image_heigth,image_width)
+    width_shift,heigth_shift = image_width//2,0
+    # cv2.imwrite("squre.jpg",img[heigth_shift:,width_shift:,:])
+    bboxes = headDetector.run(img[:,width_shift:, :], get_largest=True)  # 人脸检测，获取面积最大的人脸
+    if not bboxes:
+        return [0, 0, 0, 0], 0
+    box = [int(width_shift+bboxes[0][0]),int(heigth_shift+bboxes[0][1]),int(width_shift+bboxes[0][2]),int(heigth_shift+bboxes[0][3])]
+
+    # box[0], box[1] = max(0, box[0]), max(0, box[1])
+    # box[2], box[3] = min(w, box[2]), min(h, box[3])
+    score = bboxes[0][-1]
+
+    if (box[2] - box[0]) == 0 or (box[3] - box[1]) == 0:
+        return [0, 0, 0, 0], 0
+
+    # print(box, pred, score)
+    if show_res:
+        x0, y0, x1, y1 = box
+        print(box)
+        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), thickness=2)
+        if show_crop:
+            _, crop_box = crop_face_square_rate(img, box, rate=-0.07)
+            cx0, cy0, cx1, cy1 = crop_box
+            print(crop_box)
+            cv2.rectangle(img, (cx0, cy0), (cx1, cy1), (0, 255, 0), thickness=2)
+        cv2.imshow('res', img)
+        cv2.waitKey(0)
+
+    return box, score
+
+
+if __name__ == "__main__":
+    detect_driver_face('../../test_visual/look_down.jpg', show_res=True, show_crop=True)

head_detect/demo_ori.py

@@ -0,0 +1,58 @@
+# -*- coding:utf-8 –*-
+import os
+import sys
+
+import cv2
+
+sys.path.append(os.path.dirname(os.path.dirname(sys.path[0])))
+from head_detect.head_detector.head_detectorv2 import HeadDetector
+from head_detect.utils_quailty_assurance.utils_quailty_assurance import find_files, write_json
+from utils.images import crop_face_square_rate
+import numpy as np
+
+ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
+model_path = f"{ROOT_DIR}/models/Mona_HeadDetector_v1_straght.onnx"
+
+headDetector = HeadDetector(onnx_path=model_path)
+
+
+def detect_driver_face(img, show_res=False, show_crop=False):
+    if isinstance(img, str):
+        img = cv2.imread(img)
+    else:
+        img = img.copy()
+    h, w = img.shape[:2]
+    # origin_image = cv2.resize(img, (1280, 720))
+    # origin_image = img[:, w//2:, :]
+    bboxes = headDetector.run(img[:, w//2:, :], get_largest=True)  # 人脸检测，获取面积最大的人脸
+    if not bboxes:
+        return [0, 0, 0, 0], 0
+
+    box = [int(b) for b in bboxes[0][:4]]
+    box[0] += w // 2
+    box[2] += w // 2
+    # box[0], box[1] = max(0, box[0]), max(0, box[1])
+    # box[2], box[3] = min(w, box[2]), min(h, box[3])
+    score = bboxes[0][-1]
+
+    if (box[2] - box[0]) == 0 or (box[3] - box[1]) == 0:
+        return [0, 0, 0, 0], 0
+
+    # print(box, pred, score)
+    if show_res:
+        x0, y0, x1, y1 = box
+        print(box)
+        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 0, 255), thickness=2)
+        if show_crop:
+            _, crop_box = crop_face_square_rate(img, box, rate=-0.07)
+            cx0, cy0, cx1, cy1 = crop_box
+            print(crop_box)
+            cv2.rectangle(img, (cx0, cy0), (cx1, cy1), (0, 255, 0), thickness=2)
+        cv2.imshow('res', img)
+        cv2.waitKey(0)
+
+    return box, score
+
+
+if __name__ == "__main__":
+    detect_driver_face('../../test_visual/look_down.jpg', show_res=True, show_crop=True)

head_detect/head_detector/face_detector.py

@@ -0,0 +1,258 @@
+#针对1280X720的图片
+import os
+import cv2
+import copy
+import onnxruntime
+import numpy as np
+from face_detector.face_utils import letterbox
+#from face_utils import letterbox
+
+class FaceDetector:
+    def __init__(self, onnx_path="models/smoke_phone_mosaic_v2.onnx"):
+    
+        self.onnx_path = onnx_path
+        self.onnx_session = onnxruntime.InferenceSession(self.onnx_path)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+        self.multi_class = 3                                        #修改点1：代表输出的结果为[normal smoke,phone,drink]四类 
+ 
+    def get_output_name(self, onnx_session):
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+ 
+    def get_input_name(self, onnx_session):
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+ 
+    def get_input_feed(self, input_name, image_tensor):
+
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_tensor
+        return input_feed
+ 
+    def after_process(self,pred):
+        # 输入尺寸320,192  降8、16、32倍，对应输出尺寸为(40、20、10)
+        stride = np.array([8., 16., 32.])
+        x=[pred[0],pred[1],pred[2]]
+        # ============yolov5参数 start============
+        nc=1
+        no=16 + self.multi_class
+        nl=3
+        na=3
+        #grid=[torch.zeros(1).to(device)] * nl 
+        grid=[np.zeros(1)]*nl
+        anchor_grid=np.array([[[[[[  4.,   5.]]],
+            [[[  8.,  10.]]],
+            [[[ 13.,  16.]]]]],
+            [[[[[ 23.,  29.]]],
+            [[[ 43.,  55.]]],
+            [[[ 73., 105.]]]]],
+            [[[[[146., 217.]]],
+            [[[231., 300.]]],
+            [[[335., 433.]]]]]])
+        # ============yolov5-0.5参数 end============
+        z = [] 
+        for i in range(len(x)):
+        
+            bs,ny, nx = x[i].shape[0],x[i].shape[2] ,x[i].shape[3] 
+
+            if grid[i].shape[2:4] != x[i].shape[2:4]:
+                grid[i] = self._make_grid(nx, ny)
+
+            y = np.full_like(x[i],0)
+   
+            #y[..., [0,1,2,3,4,15]] = self.sigmoid_v(x[i][..., [0,1,2,3,4,15]])
+            y[..., [0,1,2,3,4,15,16,17,18]] = self.sigmoid_v(x[i][..., [0,1,2,3,4,15,16,17,18]])
+            # 同事sigmoid_v人脸的置信度和危险动作置信度
+
+            y[..., 5:15] = x[i][..., 5:15]
+  
+
+            y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid[i]) * stride[i]  # xy
+            y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # wh
+
+            y[..., 5:7]   = y[..., 5:7] *   anchor_grid[i] + grid[i] * stride[i] # landmark x1 y1
+            y[..., 7:9]   = y[..., 7:9] *   anchor_grid[i] + grid[i] * stride[i]# landmark x2 y2
+            y[..., 9:11]  = y[..., 9:11] *  anchor_grid[i] + grid[i] * stride[i]# landmark x3 y3
+            y[..., 11:13] = y[..., 11:13] * anchor_grid[i] + grid[i] * stride[i]# landmark x4 y4
+            y[..., 13:15] = y[..., 13:15] * anchor_grid[i] + grid[i] * stride[i]# landmark x5 y5
+
+            z.append(y.reshape((bs, -1, no)))
+        return np.concatenate(z, 1)
+
+
+    def _make_grid(self, nx, ny):
+        yv, xv = np.meshgrid(np.arange(ny), np.arange(nx),indexing = 'ij')
+        return np.stack((xv, yv), 2).reshape((1, 1, ny, nx, 2)).astype(float)
+
+    def sigmoid_v(self, array):
+        return np.reciprocal(np.exp(-array) + 1.0)
+
+    def img_process(self,orgimg,long_side=320,stride_max=32):
+
+        #orgimg=cv2.imread(img_path)
+        img0 = copy.deepcopy(orgimg)
+        h0, w0 = orgimg.shape[:2]  # orig hw
+        r = long_side/ max(h0, w0)  # resize image to img_size
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)), interpolation=interp)
+        img = letterbox(img0, new_shape=(192,320),auto=False)[0] # auto True最小矩形   False固定尺度
+        # cv2.imwrite("convert.jpg",img=img)
+        # Convert        
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).transpose(2, 0, 1).copy()  # BGR to RGB, to 3x416x416
+        img = img.astype("float32")  # uint8 to fp16/32
+        img /= 255.0  # 0 - 255 to 0.0 - 1.0
+        img = img[np.newaxis,:]
+
+        return img,orgimg
+
+    def scale_coords(self,img1_shape, coords, img0_shape, ratio_pad=None):
+        # Rescale coords (xyxy) from img1_shape to img0_shape
+        if ratio_pad is None:  # calculate from img0_shape
+            gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+            pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+        else:
+            gain = ratio_pad[0][0]
+            pad = ratio_pad[1]
+        coords[:, [0, 2, 5, 7, 9, 11, 13]] -= pad[0]  # x padding
+        coords[:, [1, 3, 6, 8, 10,12, 14]] -= pad[1]  # y padding
+
+        coords[:, [0,1,2,3,5,6,7,8,9,10,11,12,13,14]] /= gain
+
+        return coords
+
+
+    def non_max_suppression(self, boxes,confs, iou_thres=0.6):
+
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
+        order = confs.flatten().argsort()[::-1]
+        keep = []
+        while order.size > 0:
+            i = order[0]
+            keep.append(i)
+            xx1 = np.maximum(x1[i], x1[order[1:]])
+            yy1 = np.maximum(y1[i], y1[order[1:]])
+            xx2 = np.minimum(x2[i], x2[order[1:]])
+            yy2 = np.minimum(y2[i], y2[order[1:]])
+            w = np.maximum(0.0, xx2 - xx1 + 1)
+            h = np.maximum(0.0, yy2 - yy1 + 1)
+            inter = w * h
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+            inds = np.where( ovr <= iou_thres)[0]
+            order = order[inds + 1]
+
+        return boxes[keep]
+
+    def nms(self, pred, conf_thres=0.1,iou_thres=0.5):
+        xc = pred[..., 4] > conf_thres
+        pred = pred[xc]
+        pred[:, 15:] *= pred[:, 4:5]
+
+        # best class only
+        confs = np.amax(pred[:, 15:16], 1, keepdims=True)
+        pred[..., 0:4] = self.xywh2xyxy(pred[..., 0:4])
+        return self.non_max_suppression(pred, confs, iou_thres)
+
+    def xywh2xyxy(self, x):
+        # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+        y = np.zeros_like(x)
+        y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+        y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+        y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+        y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+        return y
+
+    def get_largest_face(self,pred):
+        """[获取图片中最大的人脸]
+
+        Args:
+            object ([dict]): [人脸数据]
+
+        Returns:
+            [int]: [最大人脸的坐标]
+        """
+        max_index = 0
+        max_value = 0
+        for index in range(len(pred)):
+            xmin,ymin,xmax,ymax = pred[index][:4]
+            w = xmax - xmin
+            h = ymax - ymin
+            if w*h > max_value:
+                max_value = w*h
+                max_index = index
+        return max_index
+
+    def run(self, ori_image,get_largest=True):
+        #detial_dict = {}
+
+        img,orgimg=self.img_process(ori_image,long_side=320) #[1,3,640,640]
+        #print(img.shape)
+        input_feed = self.get_input_feed(self.input_name, img)
+        pred = self.onnx_session.run(self.output_name, input_feed=input_feed)
+        #detial_dict["before_decoder"] = [i.tolist() for i in pred]
+        pred=self.after_process(pred) # torch后处理
+        #detial_dict["after_decoder"] = copy.deepcopy(pred.tolist())
+
+        pred=self.nms(pred[0],0.3,0.5)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+        pred=self.scale_coords(img.shape[2:], pred, orgimg.shape)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+
+        if get_largest and pred.shape[0]!=0 :
+            pred_index = self.get_largest_face(pred)
+            pred = pred[[pred_index]]
+        bboxes = pred[:,[0,1,2,3,15]]
+        landmarks = pred[:,5:15]
+        ## 修改点，获取危险动作的标签和置信度
+        multi_class = np.argmax(pred[:,16:],axis=1)
+        multi_conf =  np.amax(pred[:, 16:], axis=1)
+        landmarks = np.reshape(landmarks,(-1,5,2))
+        return bboxes,landmarks,multi_class,multi_conf
+
+
+if __name__ == '__main__':
+    onnxmodel = FaceDetector()
+    image_path = "/tmp-data/QACode/yolov5_quality_assurance/ONNX_Inference/smoke_phone/03ajqj5uqb013bc5_1614529990_1614530009.jpg"
+    #image = np.ones(shape=[1,3,192,320], dtype=np.float32)
+    img = cv2.imread(image_path)
+    img_resize = cv2.resize(img,(320,180))
+    img_resize = cv2.copyMakeBorder(img_resize, 6, 6, 0, 0, cv2.BORDER_CONSTANT, cv2.BORDER_CONSTANT, value=(114, 114, 114))
+    #img_resize = cv2.cvtColor(img_resize,cv2.COLOR_BGR2RGB)
+    img_resize = cv2.cvtColor(img_resize,cv2.COLOR_BGR2RGB).transpose(2, 0, 1).copy()  # BGR to RGB, and 交换 h,w c to c,h,w
+    img_resize = img_resize.astype("float32")  # uint8 to fp16/32
+    img_resize /= 255.0  # 0 - 255 to 0.0 - 1.0
+    img_resize = img_resize[np.newaxis,:] # 增加维度
+    detial_dict = {}
+    input_feed = onnxmodel.get_input_feed(onnxmodel.input_name, img_resize)
+    pred = onnxmodel.onnx_session.run(onnxmodel.output_name, input_feed=input_feed)
+    detial_dict["before_decoder"] = [i.tolist() for i in pred]
+    pred=onnxmodel.after_process(pred) # torch后处理
+    detial_dict["after_decoder"] = copy.deepcopy(pred.tolist())
+    print(pred)
+    write_json("test.json",detial_dict)
+    pred=onnxmodel.nms(pred[0],0.3,0.5)
+    #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+    #pred=onnxmodel.scale_coords(img.shape[2:], pred, img.shape)
+    detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+    write_json("test.json",detial_dict)
+
+
+
+    
+
+
+
+
+
+
+

head_detect/head_detector/face_utils.py

@@ -0,0 +1,49 @@
+import cv2
+import math 
+import numpy as np
+
+
+def make_divisible(x, divisor):
+    # Returns x evenly divisible by divisor
+    return math.ceil(x / divisor) * divisor
+
+def check_img_size(img_size, s=32):
+    # Verify img_size is a multiple of stride s
+    new_size = make_divisible(img_size, int(s))  # ceil gs-multiple
+    if new_size != img_size:
+        print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
+    return new_size
+
+
+
+def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
+    # 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 (new / old)
+    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+    if not scaleup:  # only scale down, do not scale up (for better val mAP)
+        r = min(r, 1.0)
+
+    # Compute padding
+    ratio = r, 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
+    if auto:  # minimum rectangle
+        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
+    elif scaleFill:  # stretch
+        dw, dh = 0.0, 0.0
+        new_unpad = (new_shape[1], new_shape[0])
+        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
+
+    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=color)  # add border
+    return im, ratio, (dw, dh)

head_detect/head_detector/head_detectorv2.py

@@ -0,0 +1,168 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   face_detectorv2.py
+@Time    :   2022/06/23 18:51:01
+@Author  :   Xie WenZhen 
+@Version :   1.0
+@Contact :   xiewenzhen@didiglobal.com
+@Desc    :   [人头检测去解码版本v2]
+'''
+
+# here put the import lib
+import os
+import cv2
+import copy
+import onnxruntime
+import numpy as np
+from head_detect.head_detector.face_utils import letterbox
+
+class HeadDetector:
+    def __init__(self, onnx_path="models/Mona_HeadDetector_v1_straght.onnx"):
+        self.onnx_path = onnx_path
+        self.onnx_session = onnxruntime.InferenceSession(self.onnx_path)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+
+    def get_output_name(self, onnx_session):
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+ 
+    def get_input_name(self, onnx_session):
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+ 
+    def get_input_feed(self, input_name, image_tensor):
+
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_tensor
+        return input_feed
+ 
+    def img_process(self,orgimg,long_side=320,stride_max=32):
+
+        #orgimg=cv2.imread(img_path)
+        img0 = copy.deepcopy(orgimg)
+        h0, w0 = orgimg.shape[:2]  # orig hw
+        r = long_side/ max(h0, w0)  # resize image to img_size
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            # interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            interp = cv2.INTER_LINEAR
+
+            img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)), interpolation=interp)
+        img = letterbox(img0, new_shape=(192,320),auto=False)[0] # auto True最小矩形   False固定尺度
+        # cv2.imwrite("convert.jpg",img=img)
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).transpose(2, 0, 1).copy()  # BGR to RGB, to 3x416x416
+        img = img.astype("float32")  # uint8 to fp16/32
+        img /= 255.0  # 0 - 255 to 0.0 - 1.0
+        img = img[np.newaxis,:]
+
+        return img,orgimg
+
+    def scale_coords(self,img1_shape, coords, img0_shape, ratio_pad=None):
+        # Rescale coords (xyxy) from img1_shape to img0_shape
+        if ratio_pad is None:  # calculate from img0_shape
+            gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+            pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+        else:
+            gain = ratio_pad[0][0]
+            pad = ratio_pad[1]
+        coords[:, [0, 2]] -= pad[0]  # x padding
+        coords[:, [1, 3]] -= pad[1]  # y padding
+
+        coords[:, [0,1,2,3]] /= gain
+
+        return coords
+
+
+    def non_max_suppression(self, boxes,confs, iou_thres=0.6):
+
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
+        order = confs.flatten().argsort()[::-1]
+        keep = []
+        while order.size > 0:
+            i = order[0]
+            keep.append(i)
+            xx1 = np.maximum(x1[i], x1[order[1:]])
+            yy1 = np.maximum(y1[i], y1[order[1:]])
+            xx2 = np.minimum(x2[i], x2[order[1:]])
+            yy2 = np.minimum(y2[i], y2[order[1:]])
+            w = np.maximum(0.0, xx2 - xx1 + 1)
+            h = np.maximum(0.0, yy2 - yy1 + 1)
+            inter = w * h
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+            inds = np.where( ovr <= iou_thres)[0]
+            order = order[inds + 1]
+
+        return boxes[keep]
+
+    def nms(self, pred, conf_thres=0.1,iou_thres=0.5):
+        xc = pred[..., 4] > conf_thres
+        pred = pred[xc]
+        #pred[:, 15:] *= pred[:, 4:5]
+
+        # best class only
+        confs = np.amax(pred[:, 4:5], 1, keepdims=True)
+        pred[..., 0:4] = self.xywh2xyxy(pred[..., 0:4])
+        return self.non_max_suppression(pred, confs, iou_thres)
+
+    def xywh2xyxy(self, x):
+        # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+        y = np.zeros_like(x)
+        y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+        y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+        y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+        y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+        return y
+
+    def get_largest_face(self,pred):
+        """[获取图片中最大的人脸]
+
+        Args:
+            object ([dict]): [人脸数据]
+
+        Returns:
+            [int]: [最大人脸的坐标]
+        """
+        max_index = 0
+        max_value = 0
+        for index in range(len(pred)):
+            xmin,ymin,xmax,ymax = pred[index][:4]
+            w = xmax - xmin
+            h = ymax - ymin
+            if w*h > max_value:
+                max_value = w*h
+                max_index = index
+        return max_index
+
+    def run(self, ori_image,get_largest=True):
+        img,orgimg=self.img_process(ori_image,long_side=320) #[1,3,640,640]
+        #print(img.shape)
+        input_feed = self.get_input_feed(self.input_name, img)
+        pred = self.onnx_session.run(self.output_name, input_feed=input_feed)
+
+        # pred=self.after_process(pred) # torch后处理
+        pred=self.nms(pred[0],0.3,0.5)
+        pred=self.scale_coords(img.shape[2:], pred, orgimg.shape)
+        if get_largest and pred.shape[0]!=0 :
+            pred_index = self.get_largest_face(pred)
+            pred = pred[[pred_index]]
+        bboxes = pred[:,[0,1,2,3,4]]
+        return bboxes.tolist()
+
+
+    
+
+
+
+
+
+
+

head_detect/head_detector/head_detectorv3.py

@@ -0,0 +1,216 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   face_detectorv2.py
+@Time    :   2022/06/23 18:51:01
+@Author  :   Xie WenZhen 
+@Version :   1.0
+@Contact :   xiewenzhen@didiglobal.com
+@Desc    :   [人头检测去解码版本]
+'''
+
+# here put the import lib
+import os
+import cv2
+import copy
+import onnxruntime
+import numpy as np
+from head_detect.head_detector.face_utils import letterbox
+
+class HeadDetector:
+    def __init__(self, onnx_path="models/head_models/HeadDetectorv1.1.onnx"):
+        self.onnx_path = onnx_path
+        self.onnx_session = onnxruntime.InferenceSession(self.onnx_path)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+
+    def get_output_name(self, onnx_session):
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+ 
+    def get_input_name(self, onnx_session):
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+ 
+    def get_input_feed(self, input_name, image_tensor):
+
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_tensor
+        return input_feed
+ 
+    def after_process(self,pred):
+        # 输入尺寸320,192  降8、16、32倍，对应输出尺寸为(40、20、10)
+        stride = np.array([8., 16., 32.])
+        x=[pred[0],pred[1],pred[2]]
+        # ============yolov5参数 start============
+        nl=3
+
+        #grid=[torch.zeros(1).to(device)] * nl 
+        grid=[np.zeros(1)]*nl
+        anchor_grid=np.array([[[[[[  4.,   5.]]],
+            [[[  8.,  10.]]],
+            [[[ 13.,  16.]]]]],
+            [[[[[ 23.,  29.]]],
+            [[[ 43.,  55.]]],
+            [[[ 73., 105.]]]]],
+            [[[[[146., 217.]]],
+            [[[231., 300.]]],
+            [[[335., 433.]]]]]])
+        # ============yolov5-0.5参数 end============
+        z = [] 
+        for i in range(len(x)):
+        
+            ny, nx = x[i].shape[1],x[i].shape[2]
+            if grid[i].shape[2:4] != x[i].shape[2:4]:
+                grid[i] = self._make_grid(nx, ny)
+
+            y = np.full_like(x[i],0)
+   
+            #y[..., [0,1,2,3,4,15]] = self.sigmoid_v(x[i][..., [0,1,2,3,4,15]])
+            y[..., [0,1,2,3,4]] = self.sigmoid_v(x[i][..., [0,1,2,3,4]])
+            #sigmoid_v人脸的置信度和危险动作置信度
+            y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid[i]) * stride[i]  # xy
+            y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # wh
+
+            z.append(y.reshape((1, -1, 6)))
+        return np.concatenate(z, 1)
+
+
+    def _make_grid(self, nx, ny):
+        yv, xv = np.meshgrid(np.arange(ny), np.arange(nx),indexing = 'ij')
+        return np.stack((xv, yv), 2).reshape((1, ny, nx, 2)).astype(float)
+
+    def sigmoid_v(self, array):
+        return np.reciprocal(np.exp(-array) + 1.0)
+
+    def img_process(self,orgimg,long_side=320,stride_max=32):
+
+        #orgimg=cv2.imread(img_path)
+        img0 = copy.deepcopy(orgimg)
+        h0, w0 = orgimg.shape[:2]  # orig hw
+        r = long_side/ max(h0, w0)  # resize image to img_size
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            # interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            interp = cv2.INTER_LINEAR
+
+            img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)), interpolation=interp)
+        img = letterbox(img0, new_shape=(320,320),auto=False)[0] # auto True最小矩形   False固定尺度
+        # cv2.imwrite("convert1.jpg",img=img)
+        # Convert
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).transpose(2, 0, 1).copy()  # BGR to RGB, to 3x416x416
+        img = img.astype("float32")  # uint8 to fp16/32
+        img /= 255.0  # 0 - 255 to 0.0 - 1.0
+        img = img[np.newaxis,:]
+
+        return img,orgimg
+
+    def scale_coords(self,img1_shape, coords, img0_shape, ratio_pad=None):
+        # Rescale coords (xyxy) from img1_shape to img0_shape
+        if ratio_pad is None:  # calculate from img0_shape
+            gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+            pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+        else:
+            gain = ratio_pad[0][0]
+            pad = ratio_pad[1]
+        coords[:, [0, 2]] -= pad[0]  # x padding
+        coords[:, [1, 3]] -= pad[1]  # y padding
+
+        coords[:, [0,1,2,3]] /= gain
+
+        return coords
+
+
+    def non_max_suppression(self, boxes,confs, iou_thres=0.6):
+
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
+        order = confs.flatten().argsort()[::-1]
+        keep = []
+        while order.size > 0:
+            i = order[0]
+            keep.append(i)
+            xx1 = np.maximum(x1[i], x1[order[1:]])
+            yy1 = np.maximum(y1[i], y1[order[1:]])
+            xx2 = np.minimum(x2[i], x2[order[1:]])
+            yy2 = np.minimum(y2[i], y2[order[1:]])
+            w = np.maximum(0.0, xx2 - xx1 + 1)
+            h = np.maximum(0.0, yy2 - yy1 + 1)
+            inter = w * h
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+            inds = np.where( ovr <= iou_thres)[0]
+            order = order[inds + 1]
+
+        return boxes[keep]
+
+    def nms(self, pred, conf_thres=0.1,iou_thres=0.5):
+        xc = pred[..., 4] > conf_thres
+        pred = pred[xc]
+        #pred[:, 15:] *= pred[:, 4:5]
+
+        # best class only
+        confs = np.amax(pred[:, 4:5], 1, keepdims=True)
+        pred[..., 0:4] = self.xywh2xyxy(pred[..., 0:4])
+        return self.non_max_suppression(pred, confs, iou_thres)
+
+    def xywh2xyxy(self, x):
+        # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+        y = np.zeros_like(x)
+        y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+        y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+        y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+        y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+        return y
+
+    def get_largest_face(self,pred):
+        """[获取图片中最大的人脸]
+
+        Args:
+            object ([dict]): [人脸数据]
+
+        Returns:
+            [int]: [最大人脸的坐标]
+        """
+        max_index = 0
+        max_value = 0
+        for index in range(len(pred)):
+            xmin,ymin,xmax,ymax = pred[index][:4]
+            w = xmax - xmin
+            h = ymax - ymin
+            if w*h > max_value:
+                max_value = w*h
+                max_index = index
+        return max_index
+
+    def run(self, ori_image,get_largest=True):
+        img,orgimg=self.img_process(ori_image,long_side=320) #[1,3,640,640]
+        #print(img.shape)
+        input_feed = self.get_input_feed(self.input_name, img)
+        pred = self.onnx_session.run(self.output_name, input_feed=input_feed)
+        pred=self.after_process(pred) # torch后处理
+        pred=self.nms(pred[0],0.3,0.5)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+        pred=self.scale_coords(img.shape[2:], pred, orgimg.shape)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+
+        if get_largest and pred.shape[0]!=0 :
+            pred_index = self.get_largest_face(pred)
+            pred = pred[[pred_index]]
+        bboxes = pred[:,[0,1,2,3,4]]
+        return bboxes.tolist()
+
+
+    
+
+
+
+
+
+
+

head_detect/head_detector/head_detectorv4.py

@@ -0,0 +1,216 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   face_detectorv2.py
+@Time    :   2022/06/23 18:51:01
+@Author  :   Xie WenZhen 
+@Version :   1.0
+@Contact :   xiewenzhen@didiglobal.com
+@Desc    :   [人头检测去解码版本]
+'''
+
+# here put the import lib
+import os
+import cv2
+import copy
+import onnxruntime
+import numpy as np
+from head_detect.head_detector.face_utils import letterbox
+
+class HeadDetector:
+    def __init__(self, onnx_path="models/head_models/HeadDetectorv1.3.onnx"):
+        self.onnx_path = onnx_path
+        self.onnx_session = onnxruntime.InferenceSession(self.onnx_path)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+
+    def get_output_name(self, onnx_session):
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+ 
+    def get_input_name(self, onnx_session):
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+ 
+    def get_input_feed(self, input_name, image_tensor):
+
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_tensor
+        return input_feed
+ 
+    def after_process(self,pred):
+        # 输入尺寸320,192  降8、16、32倍，对应输出尺寸为(40、20、10)
+        stride = np.array([8., 16., 32.])
+        x=[pred[0],pred[1],pred[2]]
+        # ============yolov5参数 start============
+        nl=3
+
+        #grid=[torch.zeros(1).to(device)] * nl 
+        grid=[np.zeros(1)]*nl
+        anchor_grid=np.array([[[[[[  4.,   5.]]],
+            [[[  8.,  10.]]],
+            [[[ 13.,  16.]]]]],
+            [[[[[ 23.,  29.]]],
+            [[[ 43.,  55.]]],
+            [[[ 73., 105.]]]]],
+            [[[[[146., 217.]]],
+            [[[231., 300.]]],
+            [[[335., 433.]]]]]])
+        # ============yolov5-0.5参数 end============
+        z = [] 
+        for i in range(len(x)):
+        
+            ny, nx = x[i].shape[1],x[i].shape[2]
+            if grid[i].shape[2:4] != x[i].shape[2:4]:
+                grid[i] = self._make_grid(nx, ny)
+
+            y = np.full_like(x[i],0)
+   
+            #y[..., [0,1,2,3,4,15]] = self.sigmoid_v(x[i][..., [0,1,2,3,4,15]])
+            y[..., [0,1,2,3,4]] = self.sigmoid_v(x[i][..., [0,1,2,3,4]])
+            #sigmoid_v人脸的置信度和危险动作置信度
+            y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid[i]) * stride[i]  # xy
+            y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid[i]  # wh
+
+            z.append(y.reshape((1, -1, 6)))
+        return np.concatenate(z, 1)
+
+
+    def _make_grid(self, nx, ny):
+        yv, xv = np.meshgrid(np.arange(ny), np.arange(nx),indexing = 'ij')
+        return np.stack((xv, yv), 2).reshape((1, ny, nx, 2)).astype(float)
+
+    def sigmoid_v(self, array):
+        return np.reciprocal(np.exp(-array) + 1.0)
+
+    def img_process(self,orgimg,long_side=320,stride_max=32):
+
+        #orgimg=cv2.imread(img_path)
+        img0 = copy.deepcopy(orgimg)
+        h0, w0 = orgimg.shape[:2]  # orig hw
+        r = long_side/ max(h0, w0)  # resize image to img_size
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            # interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            interp = cv2.INTER_LINEAR
+
+            img0 = cv2.resize(img0, (int(w0 * r), int(h0 * r)), interpolation=interp)
+        img = letterbox(img0, new_shape=(320,288),auto=False)[0] # auto True最小矩形   False固定尺度
+        # cv2.imwrite("convert1.jpg",img=img)
+        # Convert
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB).transpose(2, 0, 1).copy()  # BGR to RGB, to 3x416x416
+        img = img.astype("float32")  # uint8 to fp16/32
+        img /= 255.0  # 0 - 255 to 0.0 - 1.0
+        img = img[np.newaxis,:]
+
+        return img,orgimg
+
+    def scale_coords(self,img1_shape, coords, img0_shape, ratio_pad=None):
+        # Rescale coords (xyxy) from img1_shape to img0_shape
+        if ratio_pad is None:  # calculate from img0_shape
+            gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  # gain  = old / new
+            pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  # wh padding
+        else:
+            gain = ratio_pad[0][0]
+            pad = ratio_pad[1]
+        coords[:, [0, 2]] -= pad[0]  # x padding
+        coords[:, [1, 3]] -= pad[1]  # y padding
+
+        coords[:, [0,1,2,3]] /= gain
+
+        return coords
+
+
+    def non_max_suppression(self, boxes,confs, iou_thres=0.6):
+
+        x1 = boxes[:, 0]
+        y1 = boxes[:, 1]
+        x2 = boxes[:, 2]
+        y2 = boxes[:, 3]
+        areas = (x2 - x1 + 1) * (y2 - y1 + 1) 
+        order = confs.flatten().argsort()[::-1]
+        keep = []
+        while order.size > 0:
+            i = order[0]
+            keep.append(i)
+            xx1 = np.maximum(x1[i], x1[order[1:]])
+            yy1 = np.maximum(y1[i], y1[order[1:]])
+            xx2 = np.minimum(x2[i], x2[order[1:]])
+            yy2 = np.minimum(y2[i], y2[order[1:]])
+            w = np.maximum(0.0, xx2 - xx1 + 1)
+            h = np.maximum(0.0, yy2 - yy1 + 1)
+            inter = w * h
+            ovr = inter / (areas[i] + areas[order[1:]] - inter)
+            inds = np.where( ovr <= iou_thres)[0]
+            order = order[inds + 1]
+
+        return boxes[keep]
+
+    def nms(self, pred, conf_thres=0.1,iou_thres=0.5):
+        xc = pred[..., 4] > conf_thres
+        pred = pred[xc]
+        #pred[:, 15:] *= pred[:, 4:5]
+
+        # best class only
+        confs = np.amax(pred[:, 4:5], 1, keepdims=True)
+        pred[..., 0:4] = self.xywh2xyxy(pred[..., 0:4])
+        return self.non_max_suppression(pred, confs, iou_thres)
+
+    def xywh2xyxy(self, x):
+        # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+        y = np.zeros_like(x)
+        y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+        y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+        y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+        y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+        return y
+
+    def get_largest_face(self,pred):
+        """[获取图片中最大的人脸]
+
+        Args:
+            object ([dict]): [人脸数据]
+
+        Returns:
+            [int]: [最大人脸的坐标]
+        """
+        max_index = 0
+        max_value = 0
+        for index in range(len(pred)):
+            xmin,ymin,xmax,ymax = pred[index][:4]
+            w = xmax - xmin
+            h = ymax - ymin
+            if w*h > max_value:
+                max_value = w*h
+                max_index = index
+        return max_index
+
+    def run(self, ori_image,get_largest=True):
+        img,orgimg=self.img_process(ori_image,long_side=320) #[1,3,640,640]
+        #print(img.shape)
+        input_feed = self.get_input_feed(self.input_name, img)
+        pred = self.onnx_session.run(self.output_name, input_feed=input_feed)
+        pred=self.after_process(pred) # torch后处理
+        pred=self.nms(pred[0],0.3,0.5)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+        pred=self.scale_coords(img.shape[2:], pred, orgimg.shape)
+        #detial_dict["after_nms"] = copy.deepcopy(pred.tolist())
+
+        if get_largest and pred.shape[0]!=0 :
+            pred_index = self.get_largest_face(pred)
+            pred = pred[[pred_index]]
+        bboxes = pred[:,[0,1,2,3,4]]
+        return bboxes.tolist()
+
+
+    
+
+
+
+
+
+
+

head_detect/head_detector/pose.py

@@ -0,0 +1,55 @@
+
+import numpy as np
+import math
+
+class Pose:
+    def __init__(self):
+
+        pt3d = np.zeros((3, 5))
+        pt3d[0, :] = [-0.3207, 0.3101, -0.0011, -0.2578, 0.2460]
+        pt3d[1, :] = [0.2629, 0.2631, -0.0800, -0.4123, -0.4127]
+        pt3d[2, :] = [0.9560, 0.9519, 1.3194, 0.9921, 0.9899]
+        self.pt3d = pt3d * 1e5
+
+    def __call__(self, pt2d):
+
+        #pt2d = np.asarray(pt2d, np.float)
+        pt2d = np.reshape(pt2d, (5, 2)).transpose()
+        pt3d = self.pt3d
+        # 参照论文Optimum Fiducials Under Weak Perspective Projection，使用弱透视投影
+        # 减均值，排除t，便于求出R
+        pt2dm = np.zeros(pt2d.shape)
+        pt3dm = np.zeros(pt3d.shape)
+        pt2dm[0, :] = pt2d[0, :] - np.mean(pt2d[0, :])
+        pt2dm[1, :] = pt2d[1, :] - np.mean(pt2d[1, :])
+        pt3dm[0, :] = pt3d[0, :] - np.mean(pt3d[0, :])
+        pt3dm[1, :] = pt3d[1, :] - np.mean(pt3d[1, :])
+        pt3dm[2, :] = pt3d[2, :] - np.mean(pt3d[2, :])
+        # 最小二乘方法计算R
+        R1 = np.dot(np.dot(np.mat(np.dot(pt3dm, pt3dm.T)).I, pt3dm), pt2dm[0, :].T)
+        R2 = np.dot(np.dot(np.mat(np.dot(pt3dm, pt3dm.T)).I, pt3dm), pt2dm[1, :].T)
+        # 计算出f
+        f = (math.sqrt(R1[0, 0] ** 2 + R1[0, 1] ** 2 + R1[0, 2] ** 2) + math.sqrt(
+            R2[0, 0] ** 2 + R2[0, 1] ** 2 + R2[0, 2] ** 2)) / 2
+        R1 = R1 / f
+
+        R2 = R2 / f
+        R3 = np.cross(R1, R2)
+        # 使用旋转矩阵R恢复出旋转角度
+        phi = math.atan(R2[0, 2] / R3[0, 2])
+        gamma = math.atan(-R1[0, 2] / (math.sqrt(R1[0, 0] ** 2 + R1[0, 1] ** 2)))
+        theta = math.atan(R1[0, 1] / R1[0, 0])
+
+        # 使用R重新计算旋转平移矩阵，求出t
+        pt3d = np.row_stack((pt3d, np.ones((1, pt3d.shape[1]))))
+        R1_orig = np.dot(np.dot(np.mat(np.dot(pt3d, pt3d.T)).I, pt3d), pt2d[0, :].T)
+        R2_orig = np.dot(np.dot(np.mat(np.dot(pt3d, pt3d.T)).I, pt3d), pt2d[1, :].T)
+
+        t3d = np.array([R1_orig[0, 3], R2_orig[0, 3], 0]).reshape((3, 1))
+        pitch = phi * 180 / np.pi
+        yaw = gamma * 180 / np.pi
+        roll = theta * 180 / np.pi
+
+        return pitch, yaw, roll
+
+

head_detect/models/HeadDetectorv1.6.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fb959a1126c460b7becda05ba121c6ebd0b2a2fcc116e14af065aec69430068
+size 1227728

head_detect/utils_quailty_assurance/draw_tools.py

@@ -0,0 +1,253 @@
+import os
+import cv2
+import math
+import numpy as np
+import matplotlib.pyplot as plt
+
+def show_results(img, bboxe,landmark,mask_label=None,emotion_label=None):
+    h,w,c = img.shape
+    tl = 1 or round(0.002 * (h + w) / 2) + 1  # line/font thickness
+    x1,y1,x2,y2,confidence = bboxe
+    cv2.rectangle(img, (int(x1),int(y1)), (int(x2), int(y2)), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
+    clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
+
+    for i in range(5):
+        point_x = int(landmark[i][0])
+        point_y = int(landmark[i][1])
+        cv2.circle(img, (point_x, point_y), tl+1, clors[i], -1)
+
+    tf = max(tl - 1, 1)  # font thickness
+    label = str(confidence)[:5]
+    cv2.putText(img, label, (int(x2), int(y2) - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    if mask_label!=None:
+        labels_dict = {0: 'Mask', 1: 'NoMask'}
+        color_dict = {0:[0,0,255],1:[255,255,0]}
+        cv2.putText(img, labels_dict[mask_label], (int(x1), int(y1)-5), 0, tl, color_dict[mask_label], thickness=tf, lineType=cv2.LINE_AA)
+    if emotion_label!=None:
+        emotion_str = "smile:{:.2f}".format(emotion_label)
+        cv2.putText(img, emotion_str, (int(x1), int(y1)-30), 0, tl, [255,153,255], thickness=tf, lineType=cv2.LINE_AA)
+
+    return img
+
+def draw_bboxes_landmarks(img, bboxe,landmark,multi_label=None,multi_conf=None):
+    h,w,c = img.shape
+    tl = 1 or round(0.002 * (h + w) / 2) + 1  # line/font thickness
+    x1,y1,x2,y2,confidence = bboxe
+    cv2.rectangle(img, (int(x1),int(y1)), (int(x2), int(y2)), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
+    clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
+
+    for i in range(5):
+        point_x = int(landmark[i][0])
+        point_y = int(landmark[i][1])
+        cv2.circle(img, (point_x, point_y), tl+1, clors[i], -1)
+
+    tf = max(tl - 1, 1)  # font thickness
+    label = f"face:{confidence:.3f}"
+    cv2.putText(img, label, (int(x1), int(y1) - 2), 0, tl / 2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    if multi_label!=None:
+        labels_dict = {0:"normal",1:"smoke",2:"phone",3:"drink"}
+        multi_str = f"{labels_dict[multi_label]}:{multi_conf:.3f}"
+        if y1>10:
+            cv2.putText(img, multi_str, (int(x1), int(y1)-17), 0, tl/2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+        else:
+            cv2.putText(img, multi_str, (int(x1), int(y2)+5), 0, tl/2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    return img
+
+def draw_bboxes(img, bboxe,multi_label=None,multi_conf=None):
+    h,w,c = img.shape
+    tl = 1 or round(0.002 * (h + w) / 2) + 1  # line/font thickness
+    x1,y1,x2,y2,confidence = bboxe
+    cv2.rectangle(img, (int(x1),int(y1)), (int(x2), int(y2)), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
+    tf = max(tl - 1, 1)  # font thickness
+    label = f"face:{confidence:.3f}"
+    cv2.putText(img, label, (int(x1), int(y1) - 2), 0, tl / 2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    if multi_label!=None:
+        labels_dict = {0:"normal",1:"smoke",2:"phone",3:"drink"}
+        multi_str = f"{labels_dict[multi_label]}:{multi_conf:.3f}"
+        if y1>10:
+            cv2.putText(img, multi_str, (int(x1), int(y1)-17), 0, tl/2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+        else:
+            cv2.putText(img, multi_str, (int(x1), int(y2)+5), 0, tl/2, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
+    return img
+    
+def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size=80):
+    height, width = img.shape[:2]
+    tl = 1 or round(0.002 * (height + width) / 2) + 1  # line/font thickness
+    pitch = pitch * np.pi / 180
+    yaw = -(yaw * np.pi / 180)
+    roll = roll * np.pi / 180
+
+    if tdx != None and tdy != None:
+        tdx = tdx
+        tdy = tdy
+    else:
+        tdx = width / 2
+        tdy = height / 2
+
+    # X-Axis pointing to right. drawn in red
+    x1 = size * (math.cos(yaw) * math.cos(roll)) + tdx
+    y1 = size * (math.cos(pitch) * math.sin(roll) + math.cos(roll)
+                * math.sin(pitch) * math.sin(yaw)) + tdy
+
+    # Y-Axis | drawn in green
+    #        v
+    x2 = size * (-math.cos(yaw) * math.sin(roll)) + tdx
+    y2 = size * (math.cos(pitch) * math.cos(roll) - math.sin(pitch)
+                * math.sin(yaw) * math.sin(roll)) + tdy
+
+    # Z-Axis (out of the screen) drawn in blue
+    x3 = size * (math.sin(yaw)) + tdx
+    y3 = size * (-math.cos(yaw) * math.sin(pitch)) + tdy
+
+    cv2.line(img, (int(tdx), int(tdy)), (int(x1), int(y1)), (0, 0, 255), 3)
+    cv2.line(img, (int(tdx), int(tdy)), (int(x2), int(y2)), (0, 255, 0), 3)
+    cv2.line(img, (int(tdx), int(tdy)), (int(x3), int(y3)), (255, 0, 0), 2)
+
+    return img
+
+def split_num(sort_lst):
+    if not sort_lst:
+        return []
+    len_lst = len(sort_lst)
+    i = 0
+    split_lst = []
+    tmp_lst = [sort_lst[i]]
+    while True:
+        if i + 1 == len_lst:
+            break
+        next_n = sort_lst[i+1]
+        if sort_lst[i] + 1 == next_n:
+            tmp_lst.append(next_n)
+        else:
+            split_lst.append(tmp_lst)
+            tmp_lst = [next_n]
+        i += 1
+    split_lst.append(tmp_lst)
+    return split_lst
+def expand_Scope(nums,lenght):
+    if nums[0]==0:
+        start = nums[0]
+        end = nums[-1]+2
+    else:
+        start = nums[0]-1
+        end = nums[-1]+1            
+    if nums[-1]==lenght-1:
+        start = nums[0]-2
+        end = nums[-1]
+    else:
+        start = nums[0]-1
+        end = nums[-1]+1  
+    return start,end
+def draw_plot(xs, ys=None, dt=None, title ="Orign Yaw angle",c='C0', label='Yaw', time=None,y_line=None,std=None,outline_dict=None,**kwargs):
+    """ plot result of KF with color `c`, optionally displaying the variance
+    of `xs`. Returns the list of lines generated by plt.plot()"""
+
+    if ys is None and dt is not None:
+        ys = xs
+        xs = np.arange(0, len(ys) * dt, dt)
+    if ys is None:
+        ys = xs
+        xs = range(len(ys))
+
+    lines = plt.title(title)
+    lines = plt.plot(xs, ys, color=c, label=label, **kwargs)
+    if time is None:
+        return lines
+    x0 = time*dt
+    y0 = ys[time]
+    lines = plt.scatter(time*dt,ys[time],s=20,color='b')
+
+    lines = plt.plot([x0,x0],[y0,0],'r-',lw=2)
+    if y_line is None:
+        return lines 
+
+    lines = plt.axhline(y=y_line, color='g', linestyle='--')
+    y_line_list = np.full(len(ys),y_line)
+    std_top = y_line_list+std
+    std_btm = y_line_list-std
+    plt.plot(xs, std_top, linestyle=':', color='k', lw=2)
+    plt.plot(xs, std_btm, linestyle=':', color='k', lw=2)
+    up_outline_lst = outline_dict["up"]
+    down_outline_lst = outline_dict["down"]
+
+    for nums in up_outline_lst:
+        start,end = expand_Scope(nums,len(ys))
+        plt.fill_between(xs[start:end+1], std_btm[start:end+1], np.array(ys)[start:end+1], facecolor='red', alpha=0.3)
+    for nums in down_outline_lst:
+        start,end = expand_Scope(nums,len(ys))
+        plt.fill_between(xs[start:end+1], np.array(ys)[start:end+1],std_top[start:end+1], facecolor='red', alpha=0.3)    
+    
+    plt.fill_between(xs, std_btm, std_top,
+                     facecolor='yellow', alpha=0.2)
+ 
+    return lines
+
+def draw_plot_old(xs, ys=None, dt=None, title ="Orign Yaw angle",c='C0', label='Yaw', time=None,y_line=None,std=None,outline_dict=None,**kwargs):
+    """ plot result of KF with color `c`, optionally displaying the variance
+    of `xs`. Returns the list of lines generated by plt.plot()"""
+
+    if ys is None and dt is not None:
+        ys = xs
+        xs = np.arange(0, len(ys) * dt, dt)
+    if ys is None:
+        ys = xs
+        xs = range(len(ys))
+
+    lines = plt.title(title)
+    lines = plt.plot(xs, ys, color=c, label=label, **kwargs)
+    if time is None:
+        return lines
+    x0 = time*dt
+    y0 = ys[time]
+    lines = plt.scatter(time*dt,ys[time],s=20,color='b')
+
+    lines = plt.plot([x0,x0],[y0,0],'r-',lw=2)
+    if y_line is None:
+        return lines 
+
+    lines = plt.axhline(y=y_line, color='g', linestyle='--')
+    y_line_list = np.full(len(ys),y_line)
+    std_top = y_line_list+std[0]
+    std_btm = y_line_list-std[1]
+    plt.plot(xs, std_top, linestyle=':', color='k', lw=2)
+    plt.plot(xs, std_btm, linestyle=':', color='k', lw=2)
+    up_outline_lst = outline_dict["up"]
+    down_outline_lst = outline_dict["down"]
+
+    for nums in up_outline_lst:
+        start,end = expand_Scope(nums,len(ys))
+        plt.fill_between(xs[start:end+1], std_btm[start:end+1], np.array(ys)[start:end+1], facecolor='red', alpha=0.3)
+    for nums in down_outline_lst:
+        start,end = expand_Scope(nums,len(ys))
+        plt.fill_between(xs[start:end+1], np.array(ys)[start:end+1],std_top[start:end+1], facecolor='red', alpha=0.3)    
+    
+    plt.fill_between(xs, std_btm, std_top,
+                     facecolor='yellow', alpha=0.2)
+ 
+    return lines
+def draw_sticker(src, offset, pupils, landmarks,
+                 blink_thd=0.22,
+                 arrow_color=(0, 125, 255), copy=False):
+    if copy:
+        src = src.copy()
+
+    left_eye_hight = landmarks[33, 1] - landmarks[40, 1]
+    left_eye_width = landmarks[39, 0] - landmarks[35, 0]
+
+    right_eye_hight = landmarks[87, 1] - landmarks[94, 1]
+    right_eye_width = landmarks[93, 0] - landmarks[89, 0]
+
+    # for mark in landmarks.reshape(-1, 2).astype(int):
+    #     cv2.circle(src, tuple(mark), radius=1,
+    #                color=(0, 0, 255), thickness=-1)
+
+    if left_eye_hight / left_eye_width > blink_thd:
+        cv2.arrowedLine(src, tuple(pupils[0].astype(int)),
+                        tuple((offset+pupils[0]).astype(int)), arrow_color, 2)
+
+    if right_eye_hight / right_eye_width > blink_thd:
+        cv2.arrowedLine(src, tuple(pupils[1].astype(int)),
+                        tuple((offset+pupils[1]).astype(int)), arrow_color, 2)
+
+    return src
+    

head_detect/utils_quailty_assurance/metrics.py

@@ -0,0 +1,123 @@
+# Model validation metrics
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+
+def getLabel2idx(labels):
+    label2idx = dict()
+    for i in labels:
+        if i not in label2idx:
+            label2idx[i] = len(label2idx)
+    return label2idx
+
+def calculate_all_prediction(confMatrix):
+    '''
+    计算总精度：对角线上所有值除以总数
+    '''
+    total_sum = confMatrix.sum()
+    correct_sum = (np.diag(confMatrix)).sum()
+    prediction = round(100*float(correct_sum)/float(total_sum),2)
+    return prediction
+
+def calculate_label_prediction(confMatrix,labelidx):
+    '''
+    计算某一个类标预测精度：该类被预测正确的数除以该类的总数
+    '''
+    label_total_sum = confMatrix.sum(axis=0)[labelidx]
+    label_correct_sum = confMatrix[labelidx][labelidx]
+    prediction = 0
+    if label_total_sum != 0:
+        prediction = round(100*float(label_correct_sum)/float(label_total_sum),2)
+    return prediction
+
+def calculate_label_recall(confMatrix,labelidx):
+    '''
+    计算某一个类标的召回率：
+    '''
+    label_total_sum = confMatrix.sum(axis=1)[labelidx]
+    label_correct_sum = confMatrix[labelidx][labelidx]
+    recall = 0
+    if label_total_sum != 0:
+        recall = round(100*float(label_correct_sum)/float(label_total_sum),2)
+    return recall
+
+def calculate_f1(prediction,recall):
+    if (prediction+recall)==0:
+        return 0
+    return round(2*prediction*recall/(prediction+recall),2)
+
+def ap_per_class(tp, conf, pred_cls, target_cls):
+    # tp = np.squeeze(tp,axis = 1)
+    conf = np.squeeze(conf,axis=1)
+    pred_cls = np.squeeze(pred_cls,axis=1)
+    # target_cls = target_cls.reshape((-1,1))
+
+    # Sort by objectness
+    i = np.argsort(-conf)
+    tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
+    #print(f"labels face:{len(target_cls)},pred face:{len(pred_cls)}")
+
+    # Find unique classes
+    unique_classes = np.unique(target_cls)
+
+    # Create Precision-Recall curve and compute AP for each class
+    px, py = np.linspace(0, 1, 1000), []  # for plotting
+    pr_score = 0.1  # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898
+    s = [unique_classes.shape[0], tp.shape[1]]  # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)
+    ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)
+    for ci, c in enumerate(unique_classes):
+        i = pred_cls == c
+        n_l = (target_cls == c).sum()  # number of labels
+        n_p = i.sum()  # number of predictions
+
+        if n_p == 0 or n_l == 0:
+            continue
+        else:
+            # Accumulate FPs and TPs
+            fpc = (1 - tp[i]).cumsum(0)
+            tpc = tp[i].cumsum(0)
+
+            # Recall
+            recall = tpc / (n_l + 1e-16)  # recall curve
+            r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0])  # r at pr_score, negative x, xp because xp decreases
+
+            # Precision
+            precision = tpc / (tpc + fpc)  # precision curve
+            p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0])  # p at pr_score
+
+            # AP from recall-precision curve
+            for j in range(tp.shape[1]):
+                ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
+    # Compute F1 score (harmonic mean of precision and recall)
+    f1 = 2 * p * r / (p + r + 1e-16)
+    return p, r, ap, f1, unique_classes.astype('int32')
+
+
+def compute_ap(recall, precision):
+    """ Compute the average precision, given the recall and precision curves
+    # Arguments
+        recall:    The recall curve (list)
+        precision: The precision curve (list)
+    # Returns
+        Average precision, precision curve, recall curve
+    """
+
+    # Append sentinel values to beginning and end
+    mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01]))
+    mpre = np.concatenate(([1.], precision, [0.]))
+
+    # Compute the precision envelope
+    mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
+
+    # Integrate area under curve
+    method = 'interp'  # methods: 'continuous', 'interp'
+    if method == 'interp':
+        x = np.linspace(0, 1, 101)  # 101-point interp (COCO)
+        ap = np.trapz(np.interp(x, mrec, mpre), x)  # integrate
+    else:  # 'continuous'
+        i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes
+        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve
+
+    return ap, mpre, mrec
+
+

head_detect/utils_quailty_assurance/result_to_coco.py

@@ -0,0 +1,108 @@
+from email import header
+
+
+# -*- encoding: utf-8 -*-
+'''
+@File    :   result_to_coco.py
+@Time    :   2022/04/27 15:54:13
+@Author  :   Xie WenZhen 
+@Version :   1.0
+@Contact :   xiewenzhen@didiglobal.com
+@Desc    :   None
+'''
+
+# here put the import lib
+import os
+import shutil
+from tqdm import tqdm
+from copy import deepcopy
+
+def parse_label_resultv1(label_path,img_width,img_higth):
+    with open(label_path, 'r') as fr:
+        labelList = fr.readlines()
+    face_list = []
+    for label in labelList:
+        label = label.strip().split()
+        x = float(label[1])
+        y = float(label[2])
+        w = float(label[3])
+        h = float(label[4])
+        x1 = (x - w / 2) * img_width
+        y1 = (y - h / 2) * img_higth
+        x2 = (x + w / 2) * img_width
+        y2 = (y + h / 2) * img_higth
+        face_list.append(deepcopy([x1,y1,x2-x1,y2-y1]))
+    return face_list
+
+def parse_label_resultv2(label_path,img_width,img_higth):
+    with open(label_path, 'r') as fr:
+        labelList = fr.readlines()
+    face_list = []
+    mark_list = []
+    category_id_list = []
+    for label in labelList:
+        label = label.strip().split()
+        x = float(label[1])
+        y = float(label[2])
+        w = float(label[3])
+        h = float(label[4])
+        x1 = (x - w / 2) * img_width
+        y1 = (y - h / 2) * img_higth
+        x2 = (x + w / 2) * img_width
+        y2 = (y + h / 2) * img_higth
+        face_list.append(deepcopy([x1,y1,x2-x1,y2-y1]))
+        ######
+        mx0_ = float(label[5]) * img_width
+        my0_ = float(label[6]) * img_higth
+        mx1_ = float(label[7]) * img_width
+        my1_ = float(label[8]) * img_higth
+        mx2_ = float(label[9]) * img_width
+        my2_ = float(label[10]) * img_higth
+        mx3_ = float(label[11]) * img_width
+        my3_ = float(label[12]) * img_higth
+        mx4_ = float(label[13]) * img_width
+        my4_ = float(label[14]) * img_higth
+        mark_list.append(deepcopy([mx0_,my0_,mx1_,my1_,mx2_,my2_,mx3_,my3_,mx4_,my4_]))
+        #####
+        category_id = int(label[15])
+        category_id_list.append(deepcopy(category_id))
+        
+    return face_list,mark_list,category_id_list
+def generate_coco_labels(bbox,img_height,img_width,keypoint,filename,category_id):
+    images_info  = {}
+    images_info["file_name"] = filename
+    images_info['id'] = 0
+    images_info['height'] = img_height
+    images_info['width'] = img_width
+    
+
+    anno = {}
+    anno['keypoints'] = keypoint
+    anno['image_id'] = 0
+    anno['id'] = 0
+    anno['num_keypoints'] = 13  # all keypoints are labelled
+    anno['bbox'] = bbox
+    anno['iscrowd'] = 0
+    anno['area'] = anno['bbox'][2] * anno['bbox'][3]
+    anno['category_id'] = category_id
+    final_output = {"images":images_info,
+                "annotations":anno}
+    return final_output
+def get_largest_face(face_dict_list):
+    if len(face_dict_list)==1:
+        return face_dict_list[0]["bbox"],face_dict_list[0]["label"]
+    max_id = 0
+    max_area = 0
+    for idx,face_dict in enumerate(face_dict_list):
+        recent_area = face_dict["bbox"][2] * face_dict["bbox"][3]
+        if recent_area>max_area:
+            max_id = idx
+            max_area = recent_area
+    return face_dict_list[max_id]["bbox"],face_dict_list[max_id]["label"]
+
+def main():
+    print("Hello, World!")
+
+
+if __name__ == "__main__":
+    main()

head_detect/utils_quailty_assurance/utils_quailty_assurance.py

@@ -0,0 +1,79 @@
+##检查两部分数据
+import os
+import json 
+import torch
+from copy import deepcopy
+
+def write_json(json_save_path,result_dict):
+    with open(json_save_path,"w") as f:
+        f.write(json.dumps(result_dict))
+
+def read_json(json_path):
+    with open(json_path, 'r') as f:
+        result_dict = json.load(f)
+    return result_dict
+    
+def parse_label(label_path):
+    with open(label_path, 'r') as fr:
+        labelList = fr.readlines()
+    label_list = []
+    for label in labelList:
+        label = label.strip().split()
+        l = int(label[-1])
+        label_list.append(deepcopy(l))
+    return label_list 
+
+def find_files(file_path,file_type='mp4'):
+    if not os.path.exists(file_path):
+        return []
+    result = []
+    for root, dirs, files in os.walk(file_path, topdown=False):
+        for name in files:
+            if name.split(".")[-1]==file_type:
+                result.append(os.path.join(root, name))
+    return result
+def find_images(file_path):
+    if not os.path.exists(file_path):
+        return []
+    result = []
+    for root, dirs, files in os.walk(file_path, topdown=False):
+        for name in files:
+            if name.split(".")[-1] in ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng']:
+                result.append(os.path.join(root, name))
+    return result
+
+def box_iou(box1, box2):
+    def box_area(box):
+        # box = 4xn
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area1 = box_area(box1.T)
+    area2 = box_area(box2.T)
+
+    # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
+    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) -
+             torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
+    # iou = inter / (area1 + area2 - inter)
+    return inter / (area1[:, None] + area2 - inter)
+
+def xywh2xyxy(x):
+    # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
+    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
+    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
+    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
+    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
+    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
+    return y
+    
+def match_Iou(label5_list,label16_list):
+    label15_torch = torch.tensor(label5_list)
+    label16_torch = torch.tensor(label16_list)
+    bbox5 = xywh2xyxy(label15_torch[:,1:5])
+    bbox16 = xywh2xyxy(label16_torch[:,1:5])
+    label5 = label15_torch[:,0]
+    label16 = label16_torch[:,15]
+    ious, i = box_iou(bbox5, bbox16).max(1)
+    final_result = label16_torch[i]
+    final_result[:,15] = label5[i]
+    return final_result
+

head_detect/utils_quailty_assurance/video2imglist.py

@@ -0,0 +1,57 @@
+#将视频按时间抽取关键帧
+# -*- coding: utf-8 -*-
+
+import cv2
+import math
+import os
+
+
+def video2img(flv="ceshivedio.flv",rate=0.3,start=1,end=100):
+
+    list_all=[]
+    vc=cv2.VideoCapture(flv)
+    c=1
+    fps=dps(flv)
+    if rate>fps:
+       print("the fps is %s, set the rate=fps"%(fps))
+       rate=fps
+
+    if vc.isOpened():
+        rval=True
+    else:
+        rval=False
+
+    j=1.0
+    count=0.0
+    while rval:
+
+        count=fps/rate*(j-1)
+        rval,frame=vc.read()
+        if (c-1)==int(count):
+            j+=1
+            if (math.floor(c/fps))>=start and (math.floor(c/fps))<end:
+                if frame is not None:
+                   list_all.append(frame)
+        c+=1
+
+        if (math.floor(c/fps))>=end:
+            break
+    print("[ %d ] pictures from '%s' "%(len(list_all),flv))
+    vc.release()
+
+    return list_all
+    
+def dps(vedio):
+    video = cv2.VideoCapture(vedio)
+    #(major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
+    fps = video.get(cv2.CAP_PROP_FPS)
+    video.release()
+    return fps
+
+if __name__=="__main__":
+    video_path = "/tmp-data/QACode/QAMaterial/2022-02-26video/B200C视频/低头抬头+打电话+打哈欠+抽烟.mp4"
+    imglist = video2img(video_path,rate=0.1,start=0,end=100000)
+    print(len(imglist))
+    os.makedirs("tmp",exist_ok=True)
+    for idx,image in enumerate(imglist):
+        cv2.imwrite("{}/sample_{}.jpg".format("tmp",idx),image)

inference_mtl.py

@@ -0,0 +1,266 @@
+# -*- coding:utf-8 –*-
+import os
+
+import cv2
+import numpy as np
+import torch
+
+from head_detect.demo import detect_driver_face
+from models.shufflenet2_att_m import ShuffleNetV2
+from utils.images import expand_box_rate, crop_with_pad, show
+from utils.os_util import get_file_paths
+from utils.plt_util import DrawMTL
+
+project = 'dms3'
+version = 'v1.0'
+
+if version in ["v0.1", "v0.2"]:
+    input_size = [160, 160]
+    model_size = '0.5x'
+    stack_lite_head = 1
+    resume = 'data/dms3/v0.2/epoch100_11_0.910_0.865_0.983_0.984_0.964_0.923_0.952_0.831_0.801_0.948_0.948_0.833.pth'
+    onnx_export_path = 'data/dms3/v0.2/dms3_mtl_v0.2.onnx'
+    tasks = ["ems", "eye", 'mouth', 'glass', 'mask', 'smoke', 'phone', "eyelid_r", "eyelid_l", 'shift_x', 'shift_y', 'expand']
+    classes = [['normal', 'left', 'down', 'right', 'ind'], ['normal', 'close', 'ind'], ['normal', 'yawn', 'ind'],
+               ['normal', 'glass', 'ind'], ['normal', 'mask', 'ind'], ['normal', 'smoke'], ['normal', 'phone'],
+               ['distance'], ['distance'], ['distance'], ['distance'], ['distance']]
+    task_types = [0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 3]
+    num_classes = [5, 3, 3, 3, 3, 2, 2, 1, 1, 1, 1, 1]
+    reg_relative_max = [-1.] * 7 + [0.05, 0.05, 0.7, 0.7, 0.1]
+    reg_dimension = [-1.] * 7 + [3, 3, 3, 3, 2]
+    expand_rate = -0.075
+
+elif version == 'v1.0':
+    input_size = [160, 160]
+    model_size = '0.5x'
+    stack_lite_head = [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1]
+    resume = 'data/dms3/v1.0/dms3_mtl_v1.0.pth'
+    onnx_export_path = 'data/dms3/v1.0/dms3_mtl_v1.0.onnx'
+    tasks = ["ems", "eye", 'mouth', 'glass', 'mask', 'smoke', 'phone', "eyelid_r", "eyelid_l", 'shift_x', 'shift_y',
+             'expand']
+    classes = [['normal', 'left', 'down', 'right', 'ind'], ['normal', 'close', 'ind'], ['normal', 'yawn', 'ind'],
+               ['normal', 'glass', 'ind'], ['normal', 'mask', 'ind'], ['normal', 'smoke'], ['normal', 'phone'],
+               ['distance'], ['distance'], ['distance'], ['distance'], ['distance']]
+    task_types = [0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 3]
+    num_classes = [5, 3, 3, 3, 3, 2, 2, 1, 1, 1, 1, 1]
+    reg_relative_max = [-1.] * 7 + [0.05, 0.05, 0.7, 0.7, 0.1]
+    reg_dimension = [-1.] * 7 + [3, 3, 3, 3, 2]
+    expand_rate = -0.075
+
+else:
+    raise NotImplementedError
+
+# device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+device = "cpu"
+model = ShuffleNetV2(num_tasks=len(tasks), task_types=task_types, num_classes=num_classes, model_size=model_size,
+                     with_last_conv=0, stack_lite_head=stack_lite_head, lite_head_channels=-1, onnx=False)
+model.load_state_dict(torch.load(resume, map_location=device)['state_dict'])
+model.to(device)
+print(f"loading {resume}")
+model.eval()
+
+drawer = DrawMTL(project, tasks, task_types, classes)
+
+
+def onnx_export(export_name='test.onnx'):
+    model_export = ShuffleNetV2(num_tasks=len(tasks), task_types=task_types, num_classes=num_classes, model_size=model_size,
+                                with_last_conv=0, stack_lite_head=stack_lite_head, lite_head_channels=-1, onnx=True)
+    model_export.load_state_dict(torch.load(resume, map_location=device)['state_dict'])
+    model_export.eval()
+
+    example = torch.randn(1, 3, input_size[1], input_size[0])
+    torch.onnx.export(
+        model_export,  # model being run
+        example,  # model input (or a tuple for multiple inputs)
+        export_name,
+        verbose=False,
+        # store the trained parameter weights inside the model file
+        training=False,
+        input_names=['input'],
+        output_names=tasks,
+        do_constant_folding=True
+    )
+
+
+def inference_xyl(input_img, vis=False, base_box=None, dx=None, dy=None, dl=None, return_drawn=False):
+    """
+    单张推理，人头检测+多任务分类（带xy偏移）
+    :param input_img: input image path/cv_array
+    :param vis: 可视化
+    :param base_box: 基础框
+    :param dx: x方向偏移
+    :param dy: y方向偏移
+    :param return_drawn:
+    :return: preds->list [ems_pred, eye_pred, ..., dist_r, dist_l],
+             probs->list [[ems_probs], [eye_probs], ..., [1], [1]]
+    """
+    if isinstance(input_img, str) and os.path.isfile(input_img):
+        img = cv2.imread(input_img)
+        img_name = os.path.basename(input_img)
+    else:
+        img = input_img
+        img_name = 'image'
+
+    # img = cv2.resize(img, dsize=None, fx=0.5, fy=0.5)
+    if base_box is None or dx is None or dy is None or dl is None or \
+            abs(dx) >= 0.35 or abs(dy) >= 0.35 or abs(dl) >= 0.1:
+        box, score = detect_driver_face(img.copy())
+        print(box)
+        # if box == [0, 0, 0, 0]:
+        #     return
+        base_box = expand_box_rate(img, box, rate=expand_rate)
+        print(base_box)
+    else:
+        box = None
+        w, h = base_box[2] - base_box[0], base_box[3] - base_box[1]
+        assert w == h
+        x0, y0, x1, y1 = base_box
+        x0, x1 = x0 + int(w * dx), x1 + int(w * dx)
+        y0, y1 = y0 + int(h * dy), y1 + int(h * dy)
+
+        expand = int(h * dl / (1 + 2 * dl))
+        x0, y0 = x0 + expand, y0 + expand
+        x1, y1 = x1 - expand, y1 - expand
+
+        base_box = [x0, y0, x1, y1]
+
+    crop_img = crop_with_pad(img, base_box)
+    crop_img = cv2.resize(crop_img, tuple(input_size))
+    crop_img = crop_img.astype(np.float32)
+    crop_img = (crop_img - 128) / 127.
+    crop_img = crop_img.transpose([2, 0, 1])
+    crop_img = torch.from_numpy(crop_img).to(device)
+    crop_img = crop_img.view(1, *crop_img.size())
+
+    outputs = model(crop_img)
+    preds, probs = [], []
+    msg_list = []
+    for ti, outs in enumerate(outputs):
+        if task_types[ti] == 0:
+            sub_probs = torch.softmax(outs, dim=1).cpu().detach().numpy()[0]
+            sub_pred = np.argmax(sub_probs)
+            # msg_list.append(f'{tasks[ti].upper()}: {classes[ti][sub_pred]} {sub_probs[sub_pred]:.3f}')
+            msg_list.append(f'{tasks[ti].upper()}: {sub_probs}')
+            preds.append(sub_pred)
+            probs.append([round(x, 3) for x in sub_probs])
+        elif task_types[ti] == 1:
+            sub_pred = outs.cpu().detach().item() * (base_box[2] - base_box[0]) * reg_relative_max[ti] / reg_dimension[ti]
+            msg_list.append(f'{tasks[ti].upper()}: {sub_pred:.6f}')
+            preds.append(sub_pred)
+            probs.append([round(sub_pred, 3)])
+        elif task_types[ti] in [2, 3]:
+            sub_pred = outs.cpu().detach().item() * reg_relative_max[ti] / reg_dimension[ti]
+            msg_list.append(f'{tasks[ti].upper()}: {sub_pred:.6f}')
+            preds.append(sub_pred)
+            probs.append([round(sub_pred, 3)])
+
+    # print('\n'.join(msg_list))
+
+    if vis:
+        # drawn = draw_texts(img, msg_list, box, crop_box, use_mask=True)
+        # show(drawn, img_name)
+        drawn = drawer.draw_result(img, preds, probs, box, base_box)
+        # drawn = drawer.draw_ind(img)
+        show(drawn, img_name)
+
+    if return_drawn:
+        drawn = drawer.draw_result(img, preds, probs, box, base_box)
+        return drawn
+
+    return preds, probs, box, base_box, msg_list
+
+
+def inference_xyl_dir(data_dir, vis):
+    img_paths = get_file_paths(data_dir)
+    for p in img_paths:
+        print(f"\n{os.path.basename(p)}")
+        inference_xyl(p, vis=vis)
+
+
+def inference_with_basebox(input_img, base_box):
+    if isinstance(input_img, str) and os.path.isfile(input_img):
+        img = cv2.imread(input_img)
+    else:
+        img = input_img
+
+    crop_img = crop_with_pad(img, base_box)
+    crop_img = cv2.resize(crop_img, tuple(input_size))
+    crop_img = crop_img.astype(np.float32)
+    crop_img = (crop_img - 128) / 127.
+    crop_img = crop_img.transpose([2, 0, 1])
+    crop_img = torch.from_numpy(crop_img).to(device)
+    crop_img = crop_img.view(1, *crop_img.size())
+
+    outputs = model(crop_img)
+    preds, probs = [], []
+    msg_list = []
+    for ti, outs in enumerate(outputs):
+        if task_types[ti] == 0:
+            sub_probs = torch.softmax(outs, dim=1).cpu().detach().numpy()[0]
+            sub_pred = np.argmax(sub_probs)
+            # msg_list.append(f'{tasks[ti].upper()}: {classes[ti][sub_pred]} {sub_probs[sub_pred]:.3f}')
+            msg_list.append(f'{tasks[ti].upper()}: {sub_probs}')
+            preds.append(sub_pred)
+            # probs.append([round(x, 3) for x in sub_probs])
+            probs += sub_probs.tolist()
+        elif task_types[ti] == 1:
+            sub_pred = outs.cpu().detach().item() * (base_box[2] - base_box[0]) * reg_relative_max[ti] / reg_dimension[
+                ti]
+            msg_list.append(f'{tasks[ti].upper()}: {sub_pred:.6f}')
+            preds.append(sub_pred)
+            # probs.append([round(sub_pred, 3)])
+            probs.append(sub_pred)
+        elif task_types[ti] in [2, 3]:
+            sub_pred = outs.cpu().detach().item() * reg_relative_max[ti] / reg_dimension[ti]
+            msg_list.append(f'{tasks[ti].upper()}: {sub_pred:.6f}')
+            preds.append(sub_pred)
+            # probs.append([round(sub_pred, 3)])
+            probs.append(sub_pred)
+
+    return probs
+
+
+def generate_onnx_result(data_dir, label_file, save_file):
+    from utils.labels import load_labels, save_labels
+    from utils.multiprogress import MultiThreading
+
+    label_dict = load_labels(label_file)
+
+    def kernel(img_name):
+        img_path = os.path.join(data_dir, img_name)
+        box = [int(b) for b in label_dict[img_name][:4]]
+        probs = inference_with_basebox(img_path, box)
+        print(img_name, probs)
+        return img_name, probs
+
+    exe = MultiThreading(label_dict.keys(), 8)
+    res = exe.run(kernel)
+    res_dict = {r[0]: r[1] for r in res}
+    save_labels(save_file, res_dict)
+
+
+if __name__ == '__main__':
+    inference_xyl('/Users/didi/Desktop/MTL/dataset/images/20210712062707_09af469be0e74945994d0d6e9e0cbe36_209949.jpg', vis=True)
+    # inference_xyl_dir('test_mtl', vis=True)
+
+    # generate_onnx_result('/Users/didi/Desktop/MTL/dataset/images',
+    #                      '/Users/didi/Desktop/test_mtl_raw_out/test_mtl_raw_label.txt',
+    #                      '/Users/didi/Desktop/test_mtl_raw_out/test_mtl_raw_torch_prob.txt')
+
+    # onnx_export_pipeline(onnx_export_path, export_func=onnx_export, view_net=True, simplify=True)
+
+"""
+20210712062707_09af469be0e74945994d0d6e9e0cbe36_209949.jpg
+EMS: 	[1.1826814e-06 5.6046390e-09 9.9999881e-01 6.0338436e-09 5.0259811e-08]
+EYE: 	[9.999689e-01 3.101773e-05 1.121569e-07]
+MOUTH: 	[9.991779e-01 5.185943e-06 8.168342e-04]
+GLASS: 	[9.9998879e-01 1.0724665e-05 4.9968133e-07]
+MASK: 	[9.9702114e-01 2.5327259e-03 4.4608722e-04]
+SMOKE: 	[0.94698274 0.05301724]
+PHONE: 	[9.9996448e-01 3.5488123e-05]
+EYELID_R: 	[1.9863424]
+EYELID_L: 	[2.1488686]
+SHIFT_X: 	[0.00669874]
+SHIFT_Y: 	[-0.00329317]
+EXPAND: 	[-0.07628014]
+"""

inference_video_mtl.py

@@ -0,0 +1,224 @@
+# -*- coding:utf-8 –*-
+import pandas as pd
+
+from inference_mtl import *
+from utils.os_util import get_file_paths
+from utils.plt_util import plot_scores_mtl, syn_plot_scores_mtl, DrawMTL
+from utils.time_util import convert_input_time, convert_stamp
+
+
+FOURCC = 'x264'  # output video codec [h264] for internet
+
+
+def inference_videos(input_video, save_dir='output/out_vdo', detect_mode='frame', frequency=1.0, continuous=True,
+                     show_res=False, plot_score=False, save_score=False, save_vdo=False, save_img=False, img_save_dir=None,
+                     syn_plot=True, resize=None, time_delta=1, save_vdo_path=None, save_plt_path=None, save_csv_path=None):
+    """
+    @param input_video: 输入视频路径，文件夹或单个视频
+    @param save_dir: 输出视频存储目录
+    @param detect_mode: 抽帧模式，second/frame，表示按秒/帧推理
+    @param frequency: 抽帧频率倍数（0-n），乘法关系，数值越大间隔越大，可输入小数
+    @param show_res: 是否可视化推理过程的结果
+    @param continuous: 输入一个文件夹时是否其中的视频连续推理
+    @param plot_score: 是否绘制分数图
+    @param save_score: 是否保存推理结果csv文件
+    @param save_vdo: 是否保存推理结果视频
+    @param save_img: 是否保存中间帧图片，需单独设置保存条件
+    @param img_save_dir: 图像保存目录
+    @param syn_plot: 是否绘制动态同步的分数图（人数座次不支持）
+    @param resize: 输出调整
+    @param time_delta: 展示视频间隔 0 按键下一张
+    @param save_vdo_path: 视频保存路径，默认为None
+    @param save_plt_path: 分数图保存路径，默认为None
+    @param save_csv_path: 分数明细保存路径，默认为None
+    @return:
+    """
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    if save_img:
+        img_save_dir = os.path.join(save_dir, 'images') if img_save_dir is None else img_save_dir
+        if not os.path.exists(img_save_dir):
+            os.makedirs(img_save_dir)
+        save_count = 0
+
+    separately_save = True
+    if os.path.isfile(input_video):
+        vdo_list = [input_video]
+    elif os.path.isdir(input_video):
+        vdo_list = get_file_paths(input_video, mod='vdo')
+        if continuous:
+            title = os.path.basename(input_video)
+            save_vdo_path = os.path.join(save_dir, title + '.mp4') if save_vdo_path is None else save_vdo_path
+            save_plt_path = os.path.join(save_dir, title + '.jpg') if save_plt_path is None else save_plt_path
+            save_csv_path = os.path.join(save_dir, title + '.csv') if save_csv_path is None else save_csv_path
+            separately_save = False
+            frames, seconds = 0, 0
+    else:
+        print(f'No {input_video}')
+        return
+
+    if save_score:
+        columns = ['index']
+        for ti, task in enumerate(tasks):
+            if ti < 7:
+                sub_columns = [f"{task}-{sc}" for sc in classes[ti]]
+            else:
+                sub_columns = [task]
+            columns += sub_columns
+
+    if save_vdo and not separately_save:
+        video = cv2.VideoCapture(vdo_list[0])
+        width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) if resize is None else resize[0]
+        height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) if resize is None else resize[1]
+        fps = video.get(cv2.CAP_PROP_FPS)
+        fourcc = cv2.VideoWriter_fourcc(*FOURCC)
+        out_video = cv2.VideoWriter(save_vdo_path, fourcc, fps if detect_mode == 'frame' else 5,
+                                    (int(width*1.5), height) if syn_plot else (width, height))
+        print(f"result video save in '{save_vdo_path}'")
+
+    res_list = []
+    for vdo_path in vdo_list:
+        vdo_name = os.path.basename(vdo_path)
+        try:
+            start_time_str = vdo_name.split('_')[2][:14]
+            start_time_stamp = convert_input_time(start_time_str, digit=10)
+        except:
+            start_time_str, start_time_stamp = '', 0
+
+        cap = cv2.VideoCapture(vdo_path)
+        cur_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        if not cur_frames:
+            continue
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        cur_seconds = int(cur_frames / (fps + 1e-6))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        print(f"video:{vdo_name}  width:{width}  height:{height}  fps:{fps:.1f}  frames:{cur_frames}  "
+              f"seconds: {cur_seconds}  start_time:{start_time_str}")
+
+        if resize is not None:
+            width, height = resize
+
+        cur_res_list = []
+        if separately_save:
+            title = os.path.splitext(vdo_name)[0]
+            save_vdo_path = os.path.join(save_dir, title + '_res.mp4') if save_vdo_path is None else save_vdo_path
+            save_plt_path = os.path.join(save_dir, title + '_res.jpg') if save_plt_path is None else save_plt_path
+            save_csv_path = os.path.join(save_dir, title + '_res.csv') if save_csv_path is None else save_csv_path
+            if save_vdo:
+                fourcc = cv2.VideoWriter_fourcc(*FOURCC)
+                out_video = cv2.VideoWriter(save_vdo_path, fourcc, fps if detect_mode == 'frame' else 5,
+                                            (int(1.5*width), height) if syn_plot else (width, height))
+                print(f"result video save in '{save_vdo_path}'")
+        else:
+            frames += cur_frames
+            seconds += cur_seconds
+
+        base_box, dx, dy, dl = None, None, None, None
+
+        step = 1 if detect_mode == 'frame' else fps
+        step = max(1, round(step * frequency))
+        count = 0
+        for i in range(0, cur_frames, step):
+            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
+            ret, frame = cap.read()
+            if not ret:
+                # print('video end!')
+                break
+            count += 1
+
+            if resize is not None and frame.shape[0] != resize[1] and frame.shape[1] != resize[0]:
+                frame = cv2.resize(frame, resize)
+
+            cur_res = inference_xyl(frame, vis=False, base_box=base_box, dx=dx, dy=dy, dl=dl)
+            if cur_res is None:
+                preds = [0] * len(tasks)
+                probs = [[0] * len(sub_classes) for sub_classes in classes]
+                msg_list = ['no driver']
+                base_box, dx, dy, dl = None, None, None, None
+            else:
+                preds, probs, box, crop_box, msg_list = cur_res
+                base_box, dx, dy, dl = crop_box.copy(), preds[-3], preds[-2], preds[-1]
+
+            if start_time_stamp:
+                time_stamp = start_time_stamp + round(i/fps)
+                cur_time = convert_stamp(time_stamp)
+                cur_res_list.append(tuple([cur_time+f'-{i}'] + [round(p, 3) for sub_probs in probs for p in sub_probs]))
+                res_list.append(tuple([cur_time+f'-{i}'] + [round(p, 3) for sub_probs in probs for p in sub_probs]))
+            else:
+                cur_time = ''
+                cur_res_list.append(tuple([count] + [round(p, 3) for sub_probs in probs for p in sub_probs]))
+                res_list.append(tuple([count] + [round(p, 3) for sub_probs in probs for p in sub_probs]))
+
+            if not count % 10 and count:
+                msg = "{} {} => {}".format(i, cur_time, '\t'.join(msg_list))
+                print(msg)
+
+            if save_img and probs[1][1] > 0.8:  # Todo：设置不同的保存条件
+                img_name = vdo_name.replace(".mp4", f'_{i}.jpg') if not cur_time else \
+                    f"{convert_stamp(convert_input_time(cur_time), '%Y%m%d%H%M%S')}_{i}.jpg"
+                img_save_path = os.path.join(img_save_dir, img_name)
+                cv2.imwrite(img_save_path, frame)
+                save_count += 1
+
+            if show_res or save_vdo:
+                drawn = drawer.draw_ind(frame) if cur_res is None else \
+                    drawer.draw_result(frame, preds, probs, box, crop_box, use_mask=False, use_frame=False)
+                if syn_plot:
+                    score_array = np.array([r[1:] for r in res_list])
+                    if detect_mode == 'second' and cur_time:
+                        indexes = [r[0][-5:] for r in res_list]
+                    else:
+                        indexes = list(range(len(res_list)))
+                    window_length = 300 if detect_mode == 'frame' else 30
+                    assert len(score_array) == len(indexes)
+                    score_chart = syn_plot_scores_mtl(
+                        tasks, [[0, 1, 2, 3, 4, 5, 6], [7, 8], [9, 10], [11]], classes, indexes, score_array,
+                        int(0.5*width), height, window_length, width/1280)
+
+                    drawn = np.concatenate([drawn, score_chart], axis=1)
+
+            if show_res:
+                cv2.namedWindow(title, 0)
+                # cv2.moveWindow(title, 0, 0)
+                # cv2.setWindowProperty(title, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
+                cv2.imshow(title, drawn)
+                cv2.waitKey(time_delta)
+
+            # write the frame after processing
+            if save_vdo:
+                out_video.write(drawn)
+
+        if separately_save:
+            if show_res:
+                cv2.destroyWindow(title)
+            if plot_score:
+                res = np.array([r[1:] for r in cur_res_list])
+                plot_scores_mtl(tasks, task_types, classes, res, title, detect_mode, save_dir=save_dir,
+                                save_path=save_plt_path, show=show_res)
+            if save_score:
+                df = pd.DataFrame(cur_res_list, columns=columns)
+                df.to_csv(save_csv_path, index=False, float_format='%.3f')
+
+    if save_img:
+        print(f"total save {save_count} images")
+
+    if not separately_save:
+        if show_res:
+            cv2.destroyWindow(title)
+        if plot_score:
+            res = np.array([r[1:] for r in res_list])
+            plot_scores_mtl(tasks, task_types, classes, res, title, detect_mode, save_dir=save_dir,
+                            save_path=save_plt_path, show=show_res)
+        if save_score:
+            df = pd.DataFrame(res_list, columns=columns)
+            df.to_csv(save_csv_path, index=False, float_format='%.3f')
+
+    return save_vdo_path, save_plt_path, save_vdo_path
+
+
+if __name__ == '__main__':
+    inference_videos('/Users/didi/Desktop/CARVIDEO_03afqwj7uw5d801e_20230708132305000_20230708132325000.mp4',
+                     save_dir='/Users/didi/Desktop/error_res_v1.0',
+                     detect_mode='second', frequency=0.2, plot_score=True, save_score=True, syn_plot=True,
+                     save_vdo=True, save_img=False, continuous=False, show_res=True, resize=(1280, 720), time_delta=1)

models/module/activation.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+
+activations = {'ReLU': nn.ReLU,
+               'LeakyReLU': nn.LeakyReLU,
+               'ReLU6': nn.ReLU6,
+               'SELU': nn.SELU,
+               'ELU': nn.ELU,
+               None: nn.Identity
+               }
+
+
+def act_layers(name):
+    assert name in activations.keys()
+    if name == 'LeakyReLU':
+        return nn.LeakyReLU(negative_slope=0.1, inplace=True)
+    else:
+        return activations[name](inplace=True)

models/module/blocks.py

@@ -0,0 +1,300 @@
+# -*- coding:utf-8 –*-
+import torch.nn as nn
+from models.attention.attention_blocks import *
+import torch.nn.functional as F
+from utils.common import log_warn
+
+
+# 在每个shuffle_block之后加入att_block
+class ShuffleV2Block(nn.Module):
+    def __init__(self, inp, oup, mid_channels, ksize, stride, attention='', ratio=16, loc='side', onnx=False):
+        super(ShuffleV2Block, self).__init__()
+        self.onnx = onnx
+        self.stride = stride
+        assert stride in [1, 2]
+
+        self.mid_channels = mid_channels
+        self.ksize = ksize
+        pad = ksize // 2
+        self.pad = pad
+        self.inp = inp
+
+        outputs = oup - inp
+
+        branch_main = [
+            # pw
+            nn.Conv2d(inp, mid_channels, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            # dw
+            nn.Conv2d(mid_channels, mid_channels, ksize, stride, pad, groups=mid_channels, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            # pw-linear
+            nn.Conv2d(mid_channels, outputs, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(outputs),
+            nn.ReLU(inplace=True),
+        ]
+        self.branch_main = nn.Sequential(*branch_main)
+
+        if stride == 2:
+            branch_proj = [
+                # dw
+                nn.Conv2d(inp, inp, ksize, stride, pad, groups=inp, bias=False),
+                nn.BatchNorm2d(inp),
+                # pw-linear
+                nn.Conv2d(inp, inp, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(inp),
+                nn.ReLU(inplace=True),
+            ]
+            self.branch_proj = nn.Sequential(*branch_proj)
+        else:
+            self.branch_proj = None
+
+        if attention:
+            self.loc = loc
+            att_out = outputs if loc == 'side' else oup
+            if attention.lower() == 'se':
+                self.att_block = SELayer(att_out, reduction=ratio)
+            elif attention.lower() == 'cbam':
+                self.att_block = CBAM(att_out, ratio)
+            elif attention.lower() == 'gc':
+                self.att_block = GCBlock(att_out, ratio=ratio)
+            else:
+                raise NotImplementedError
+        else:
+            self.att_block = None
+
+    def forward(self, old_x):
+        if self.stride == 1:
+            x_proj, x = self.channel_shuffle(old_x)
+        else:
+            x_proj = old_x
+            x_proj = self.branch_proj(x_proj)
+            x = old_x
+        x = self.branch_main(x)
+        if self.att_block and self.loc == 'side':
+            x = self.att_block(x)
+        x = torch.cat((x_proj, x), 1)
+        if self.att_block and self.loc == 'after':
+            x = self.att_block(x)
+        return x
+
+    def channel_shuffle(self, x):
+        batchsize, num_channels, height, width = x.data.size()
+        if self.onnx:
+            # 由于需要将onnx模型转换为ifx模型，ifx引擎以nchw(n=1)的格式存储数据，因此做shape变换时，尽量保证按nchw(n=1)来操作
+            x = x.reshape(1, batchsize * num_channels // 2, 2, height * width)
+            x = x.permute(0, 2, 1, 3)
+            z = num_channels // 2
+            x = x.reshape(1, -1, height, width)
+            # split时避免使用x[0]、x[1]的操作，尽量使用torch的算子来实现
+            x1, x2 = torch.split(x, split_size_or_sections=[z, z], dim=1)
+            return x1, x2
+        else:
+            x = x.reshape(batchsize * num_channels // 2, 2, height * width)
+            x = x.permute(1, 0, 2)
+            x = x.reshape(2, -1, num_channels // 2, height, width)
+            return x[0], x[1]
+
+
+# 在每个shuffle_block之后加入att_block
+class QuantizableShuffleV2Block(ShuffleV2Block):
+
+    def __init__(self, *args, **kwargs):
+        if kwargs.get('attention', ''):
+            log_warn('Quantizable model not support attention blocks')
+            kwargs['attention'] = ''
+        super(QuantizableShuffleV2Block, self).__init__(*args, **kwargs)
+        self.quantized_funcs = nn.quantized.FloatFunctional()
+
+    def forward(self, old_x):
+        if self.branch_proj is None:
+            x_proj, x = self.channel_shuffle(old_x)
+        else:
+            x_proj = old_x
+            x_proj = self.branch_proj(x_proj)
+            x = old_x
+        x = self.branch_main(x)
+        x = self.quantized_funcs.cat((x_proj, x), 1)
+        return x
+
+    def channel_shuffle(self, x):
+        batchsize, num_channels, height, width = x.data.size()
+        x = x.reshape(batchsize * num_channels // 2, 2, height * width)
+        x = x.permute(1, 0, 2)
+        x = x.reshape(2, -1, num_channels // 2, height, width)
+        return x[0], x[1]
+
+
+class ShuffleV2BlockSK(nn.Module):
+    def __init__(self, inp, oup, mid_channels, *, ksize, stride):
+        super(ShuffleV2BlockSK, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        self.mid_channels = mid_channels
+        self.ksize = ksize
+        pad = ksize // 2
+        self.pad = pad
+        self.inp = inp
+
+        outputs = oup - inp
+
+        branch_main = [
+            # pw
+            nn.Conv2d(inp, mid_channels, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            # dw
+            # nn.Conv2d(mid_channels, mid_channels, ksize, stride, pad, groups=mid_channels, bias=False),
+            SKConv(mid_channels, 2, mid_channels, stride=stride, use_relu=False),
+            # SKConv2(mid_channels, 2, mid_channels, 4, stride=stride),
+            nn.BatchNorm2d(mid_channels),
+            # pw-linear
+            nn.Conv2d(mid_channels, outputs, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(outputs),
+            nn.ReLU(inplace=True),
+        ]
+        self.branch_main = nn.Sequential(*branch_main)
+
+        if stride == 2:
+            branch_proj = [
+                # dw
+                nn.Conv2d(inp, inp, ksize, stride, pad, groups=inp, bias=False),
+                # SKConv(inp, 2, inp, stride=stride, use_relu=False),
+                # SKConv2(inp, 2, inp, 4, stride=stride),
+                nn.BatchNorm2d(inp),
+                # pw-linear
+                nn.Conv2d(inp, inp, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(inp),
+                nn.ReLU(inplace=True),
+            ]
+            self.branch_proj = nn.Sequential(*branch_proj)
+        else:
+            self.branch_proj = None
+
+    def forward(self, old_x):
+        if self.stride == 1:
+            x_proj, x = self.channel_shuffle(old_x)
+            return torch.cat((x_proj, self.branch_main(x)), 1)
+        elif self.stride == 2:
+            x_proj = old_x
+            x = old_x
+            return torch.cat((self.branch_proj(x_proj), self.branch_main(x)), 1)
+
+    def channel_shuffle(self, x):
+        batchsize, num_channels, height, width = x.data.size()
+        assert (num_channels % 4 == 0)
+        x = x.reshape(batchsize * num_channels // 2, 2, height * width)
+        x = x.permute(1, 0, 2)
+        x = x.reshape(2, -1, num_channels // 2, height, width)
+        return x[0], x[1]
+
+
+class SnetBlock(ShuffleV2Block):
+    """
+    自定义了shuffle函数，其它都一样
+    """
+
+    def channel_shuffle(self, x):
+        g = 2
+        x = x.reshape(x.shape[0], g, x.shape[1] // g, x.shape[2], x.shape[3])
+        x = x.permute(0, 2, 1, 3, 4)
+        x = x.reshape(x.shape[0], -1, x.shape[3], x.shape[4])
+        x_proj = x[:, :(x.shape[1] // 2), :, :]
+        x = x[:, (x.shape[1] // 2):, :, :]
+        return x_proj, x
+
+
+def conv_bn_relu(inp, oup, kernel_size, stride, pad):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, kernel_size, stride, pad, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU(inplace=True)
+    )
+
+
+class CEM(nn.Module):
+    """
+    Context Enhancement Module
+    改进的FPN结构，c4、c5、glb分别1x1卷积(c5上次样，glb传播)，结果聚合
+    支持feat_stride 8、16
+    TODO: 最后再做一次卷积
+    """
+
+    def __init__(self, in_channels1, in_channels2, in_channels3,
+                 feat_stride=16, squeeze_channels=245, use_relu=False):
+        super(CEM, self).__init__()
+        self.feat_stride = feat_stride
+        assert feat_stride in [8, 16], f"{feat_stride} not support, select in [8, 16]"
+        if feat_stride == 8:
+            self.conv3 = nn.Conv2d(in_channels1 // 2, squeeze_channels, 1, bias=True)
+        self.conv4 = nn.Conv2d(in_channels1, squeeze_channels, 1, bias=True)
+        self.conv5 = nn.Conv2d(in_channels2, squeeze_channels, 1, bias=True)
+        self.conv_last = nn.Conv2d(in_channels3, squeeze_channels, 1, bias=True)
+        self.use_relu = use_relu
+        if use_relu:
+            self.relu7 = nn.ReLU(inplace=True)
+
+    def forward(self, inputs):
+        if self.feat_stride == 8:
+            c3_lat = self.conv3(inputs[0])
+            c4_lat = self.conv4(inputs[1])
+            c4_lat = F.interpolate(c4_lat, size=[c3_lat.size(2), c3_lat.size(3)], mode="nearest")
+            c5_lat = self.conv5(inputs[2])
+            c5_lat = F.interpolate(c5_lat, size=[c3_lat.size(2), c3_lat.size(3)], mode="nearest")
+            glb_lat = self.conv_last(inputs[3])
+            out = c3_lat + c4_lat + c5_lat + glb_lat
+        else:
+            c4_lat = self.conv4(inputs[0])
+            c5_lat = self.conv5(inputs[1])
+            c5_lat = F.interpolate(c5_lat, size=[c4_lat.size(2), c4_lat.size(3)], mode="nearest")  # 上采样
+            glb_lat = self.conv_last(inputs[2])
+            out = c4_lat + c5_lat + glb_lat
+
+        if self.use_relu:
+            out = self.relu7(out)
+        return out
+
+
+class CEM_a(nn.Module):
+    """
+    Context Enhancement Module
+    改进的FPN结构，c4、c5、glb分别1x1卷积(c5上次样，glb传播)，结果聚合
+    支持feat_stride 8、16
+    TODO: 最后再做一次卷积
+    """
+
+    def __init__(self, in_channels1, in_channels2, in_channels3,
+                 feat_stride=16, squeeze_channels=245, use_relu=False):
+        super(CEM_a, self).__init__()
+        self.feat_stride = feat_stride
+        assert feat_stride in [8, 16], f"{feat_stride} not support, select in [8, 16]"
+        if feat_stride == 8:
+            self.conv3 = nn.Conv2d(in_channels1 // 2, squeeze_channels, 1, bias=True)
+        self.conv4 = nn.Conv2d(in_channels1, squeeze_channels, 1, bias=True)
+        self.conv5 = nn.Conv2d(in_channels2, squeeze_channels, 1, bias=True)
+        self.conv_last = nn.Conv2d(in_channels3, squeeze_channels, 1, bias=True)
+        self.use_relu = use_relu
+        if use_relu:
+            self.relu7 = nn.ReLU(inplace=True)
+
+    def forward(self, inputs):
+        if self.feat_stride == 8:
+            c3_lat = self.conv3(inputs[0])
+            c4_lat = self.conv4(inputs[1])
+            c4_lat = F.interpolate(c4_lat, size=[c3_lat.size(2), c3_lat.size(3)], mode="nearest")
+            c5_lat = self.conv5(inputs[2])
+            c5_lat = F.interpolate(c5_lat, size=[c3_lat.size(2), c3_lat.size(3)], mode="nearest")
+            glb_lat = self.conv_last(inputs[3])
+            out = c3_lat + c4_lat + c5_lat + glb_lat
+        else:
+            c4_lat = self.conv4(inputs[0])
+            c5_lat = self.conv5(inputs[1])
+            c5_lat = F.interpolate(c5_lat, size=[c4_lat.size(2), c4_lat.size(3)], mode="nearest")  # 上采样
+            glb_lat = self.conv_last(inputs[2])
+            out = c4_lat + c5_lat + glb_lat
+
+        if self.use_relu:
+            out = self.relu7(out)
+        return out

models/module/conv.py

@@ -0,0 +1,340 @@
+"""
+ConvModule refers from MMDetection
+RepVGGConvModule refers from RepVGG: Making VGG-style ConvNets Great Again
+"""
+import warnings
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from models.module.activation import act_layers
+from models.module.init_weights import kaiming_init, constant_init
+from models.module.norm import build_norm_layer
+
+
+class ConvModule(nn.Module):
+    """A conv block that contains conv/norm/activation layers.
+
+    Args:
+        in_channels (int): Same as nn.Conv2d.
+        out_channels (int): Same as nn.Conv2d.
+        kernel_size (int or tuple[int]): Same as nn.Conv2d.
+        stride (int or tuple[int]): Same as nn.Conv2d.
+        padding (int or tuple[int]): Same as nn.Conv2d.
+        dilation (int or tuple[int]): Same as nn.Conv2d.
+        groups (int): Same as nn.Conv2d.
+        bias (bool or str): If specified as `auto`, it will be decided by the
+            norm_cfg. Bias will be set as True if norm_cfg is None, otherwise
+            False.
+        conv_cfg (dict): Config dict for convolution layer.
+        norm_cfg (dict): Config dict for normalization layer.
+        activation (str): activation layer, "ReLU" by default.
+        inplace (bool): Whether to use inplace mode for activation.
+        order (tuple[str]): The order of conv/norm/activation layers. It is a
+            sequence of "conv", "norm" and "act". Examples are
+            ("conv", "norm", "act") and ("act", "conv", "norm").
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias='auto',
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 activation='ReLU',
+                 inplace=True,
+                 order=('conv', 'norm', 'act')):
+        super(ConvModule, self).__init__()
+        assert conv_cfg is None or isinstance(conv_cfg, dict)
+        assert norm_cfg is None or isinstance(norm_cfg, dict)
+        assert activation is None or isinstance(activation, str)
+        self.conv_cfg = conv_cfg
+        self.norm_cfg = norm_cfg
+        self.activation = activation
+        self.inplace = inplace
+        self.order = order
+        assert isinstance(self.order, tuple) and len(self.order) == 3
+        assert set(order) == set(['conv', 'norm', 'act'])
+
+        self.with_norm = norm_cfg is not None
+        # if the conv layer is before a norm layer, bias is unnecessary.
+        if bias == 'auto':
+            bias = False if self.with_norm else True
+        self.with_bias = bias
+
+        if self.with_norm and self.with_bias:
+            warnings.warn('ConvModule has norm and bias at the same time')
+
+        # build convolution layer
+        self.conv = nn.Conv2d(  #
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias)
+        # export the attributes of self.conv to a higher level for convenience
+        self.in_channels = self.conv.in_channels
+        self.out_channels = self.conv.out_channels
+        self.kernel_size = self.conv.kernel_size
+        self.stride = self.conv.stride
+        self.padding = self.conv.padding
+        self.dilation = self.conv.dilation
+        self.transposed = self.conv.transposed
+        self.output_padding = self.conv.output_padding
+        self.groups = self.conv.groups
+
+        # build normalization layers
+        if self.with_norm:
+            # norm layer is after conv layer
+            if order.index('norm') > order.index('conv'):
+                norm_channels = out_channels
+            else:
+                norm_channels = in_channels
+            self.norm_name, norm = build_norm_layer(norm_cfg, norm_channels)
+            self.add_module(self.norm_name, norm)
+
+        # build activation layer
+        if self.activation:
+            self.act = act_layers(self.activation)
+
+        # Use msra init by default
+        self.init_weights()
+
+    @property
+    def norm(self):
+        return getattr(self, self.norm_name)
+
+    def init_weights(self):
+        if self.activation == 'LeakyReLU':
+            nonlinearity = 'leaky_relu'
+        else:
+            nonlinearity = 'relu'
+        kaiming_init(self.conv, nonlinearity=nonlinearity)
+        if self.with_norm:
+            constant_init(self.norm, 1, bias=0)
+
+    def forward(self, x, norm=True):
+        for layer in self.order:
+            if layer == 'conv':
+                x = self.conv(x)
+            elif layer == 'norm' and norm and self.with_norm:
+                x = self.norm(x)
+            elif layer == 'act' and self.activation:
+                x = self.act(x)
+        return x
+
+
+class DepthwiseConvModule(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 bias='auto',
+                 norm_cfg=dict(type='BN'),
+                 activation='ReLU',
+                 inplace=True,
+                 order=('depthwise', 'dwnorm', 'act', 'pointwise', 'pwnorm', 'act')):
+        super(DepthwiseConvModule, self).__init__()
+        assert activation is None or isinstance(activation, str)
+        self.activation = activation
+        self.inplace = inplace
+        self.order = order
+        assert isinstance(self.order, tuple) and len(self.order) == 6
+        assert set(order) == set(['depthwise', 'dwnorm', 'act', 'pointwise', 'pwnorm', 'act'])
+
+        self.with_norm = norm_cfg is not None
+        # if the conv layer is before a norm layer, bias is unnecessary.
+        if bias == 'auto':
+            bias = False if self.with_norm else True
+        self.with_bias = bias
+
+        if self.with_norm and self.with_bias:
+            warnings.warn('ConvModule has norm and bias at the same time')
+
+        # build convolution layer
+        self.depthwise = nn.Conv2d(in_channels,
+                                   in_channels,
+                                   kernel_size,
+                                   stride=stride,
+                                   padding=padding,
+                                   dilation=dilation,
+                                   groups=in_channels,
+                                   bias=bias)
+        self.pointwise = nn.Conv2d(in_channels,
+                                   out_channels,
+                                   kernel_size=1,
+                                   stride=1,
+                                   padding=0,
+                                   bias=bias)
+
+        # export the attributes of self.conv to a higher level for convenience
+        self.in_channels = self.depthwise.in_channels
+        self.out_channels = self.pointwise.out_channels
+        self.kernel_size = self.depthwise.kernel_size
+        self.stride = self.depthwise.stride
+        self.padding = self.depthwise.padding
+        self.dilation = self.depthwise.dilation
+        self.transposed = self.depthwise.transposed
+        self.output_padding = self.depthwise.output_padding
+
+        # build normalization layers
+        if self.with_norm:
+            # norm layer is after conv layer
+            _, self.dwnorm = build_norm_layer(norm_cfg, in_channels)
+            _, self.pwnorm = build_norm_layer(norm_cfg, out_channels)
+
+        # build activation layer
+        if self.activation:
+            self.act = act_layers(self.activation)
+
+        # Use msra init by default
+        self.init_weights()
+
+    def init_weights(self):
+        if self.activation == 'LeakyReLU':
+            nonlinearity = 'leaky_relu'
+        else:
+            nonlinearity = 'relu'
+        kaiming_init(self.depthwise, nonlinearity=nonlinearity)
+        kaiming_init(self.pointwise, nonlinearity=nonlinearity)
+        if self.with_norm:
+            constant_init(self.dwnorm, 1, bias=0)
+            constant_init(self.pwnorm, 1, bias=0)
+
+    def forward(self, x, norm=True):
+        for layer_name in self.order:
+            if layer_name != 'act':
+                layer = self.__getattr__(layer_name)
+                x = layer(x)
+            elif layer_name == 'act' and self.activation:
+                x = self.act(x)
+        return x
+
+
+class RepVGGConvModule(nn.Module):
+    """
+    RepVGG Conv Block from paper RepVGG: Making VGG-style ConvNets Great Again
+    https://arxiv.org/abs/2101.03697
+    https://github.com/DingXiaoH/RepVGG
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 activation='ReLU',
+                 padding_mode='zeros',
+                 deploy=False):
+        super(RepVGGConvModule, self).__init__()
+        assert activation is None or isinstance(activation, str)
+        self.activation = activation
+
+        self.deploy = deploy
+        self.groups = groups
+        self.in_channels = in_channels
+
+        assert kernel_size == 3
+        assert padding == 1
+
+        padding_11 = padding - kernel_size // 2
+
+        # build activation layer
+        if self.activation:
+            self.act = act_layers(self.activation)
+
+        if deploy:
+            self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                                         kernel_size=kernel_size, stride=stride,
+                                         padding=padding, dilation=dilation, groups=groups, bias=True,
+                                         padding_mode=padding_mode)
+
+        else:
+            self.rbr_identity = nn.BatchNorm2d(
+                num_features=in_channels) if out_channels == in_channels and stride == 1 else None
+
+            self.rbr_dense = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                                                     kernel_size=kernel_size, stride=stride, padding=padding,
+                                                     groups=groups, bias=False),
+                                           nn.BatchNorm2d(num_features=out_channels))
+
+            self.rbr_1x1 = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                                                   kernel_size=1, stride=stride, padding=padding_11,
+                                                   groups=groups, bias=False),
+                                         nn.BatchNorm2d(num_features=out_channels))
+            print('RepVGG Block, identity = ', self.rbr_identity)
+
+    def forward(self, inputs):
+        if hasattr(self, 'rbr_reparam'):
+            return self.act(self.rbr_reparam(inputs))
+
+        if self.rbr_identity is None:
+            id_out = 0
+        else:
+            id_out = self.rbr_identity(inputs)
+
+        return self.act(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)
+
+    #   This func derives the equivalent kernel and bias in a DIFFERENTIABLE way.
+    #   You can get the equivalent kernel and bias at any time and do whatever you want,
+    #   for example, apply some penalties or constraints during training, just like you do to the other models.
+    #   May be useful for quantization or pruning.
+    def get_equivalent_kernel_bias(self):
+        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
+        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
+        kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
+        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
+
+    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
+        if kernel1x1 is None:
+            return 0
+        else:
+            return nn.functional.pad(kernel1x1, [1, 1, 1, 1])
+
+    def _fuse_bn_tensor(self, branch):
+        if branch is None:
+            return 0, 0
+        if isinstance(branch, nn.Sequential):
+            kernel = branch[0].weight
+            running_mean = branch[1].running_mean
+            running_var = branch[1].running_var
+            gamma = branch[1].weight
+            beta = branch[1].bias
+            eps = branch[1].eps
+        else:
+            assert isinstance(branch, nn.BatchNorm2d)
+            if not hasattr(self, 'id_tensor'):
+                input_dim = self.in_channels // self.groups
+                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
+                for i in range(self.in_channels):
+                    kernel_value[i, i % input_dim, 1, 1] = 1
+                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
+            kernel = self.id_tensor
+            running_mean = branch.running_mean
+            running_var = branch.running_var
+            gamma = branch.weight
+            beta = branch.bias
+            eps = branch.eps
+        std = (running_var + eps).sqrt()
+        t = (gamma / std).reshape(-1, 1, 1, 1)
+        return kernel * t, beta - running_mean * gamma / std
+
+    def repvgg_convert(self):
+        kernel, bias = self.get_equivalent_kernel_bias()
+        return kernel.detach().cpu().numpy(), bias.detach().cpu().numpy(),

models/module/fpn.py

@@ -0,0 +1,165 @@
+# -*- coding:utf-8 –*-
+"""
+from MMDetection
+"""
+
+import torch.nn as nn
+import torch.nn.functional as F
+from models.module.conv import ConvModule
+from models.module.init_weights import xavier_init
+
+
+class FPN(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 num_outs,
+                 start_level=0,
+                 end_level=-1,
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 activation=None
+                 ):
+        super(FPN, self).__init__()
+        assert isinstance(in_channels, list)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.num_ins = len(in_channels)
+        self.num_outs = num_outs
+        self.fp16_enabled = False
+
+        if end_level == -1:
+            self.backbone_end_level = self.num_ins
+            assert num_outs >= self.num_ins - start_level
+        else:
+            # if end_level < inputs, no extra level is allowed
+            self.backbone_end_level = end_level
+            assert end_level <= len(in_channels)
+            assert num_outs == end_level - start_level
+        self.start_level = start_level
+        self.end_level = end_level
+        self.lateral_convs = nn.ModuleList()
+
+        for i in range(self.start_level, self.backbone_end_level):
+            l_conv = ConvModule(
+                in_channels[i],
+                out_channels,
+                1,
+                conv_cfg=conv_cfg,
+                norm_cfg=norm_cfg,
+                activation=activation,
+                inplace=False)
+
+            self.lateral_convs.append(l_conv)
+        self.init_weights()
+
+    # default init_weights for conv(msra) and norm in ConvModule
+    def init_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                xavier_init(m, distribution='uniform')
+
+    def forward(self, inputs):
+        assert len(inputs) == len(self.in_channels)
+
+        # build laterals
+        laterals = [
+            lateral_conv(inputs[i + self.start_level])
+            for i, lateral_conv in enumerate(self.lateral_convs)
+        ]
+
+        # build top-down path
+        used_backbone_levels = len(laterals)
+        for i in range(used_backbone_levels - 1, 0, -1):
+            prev_shape = laterals[i - 1].shape[2:]
+            laterals[i - 1] += F.interpolate(
+                laterals[i], size=prev_shape, mode='bilinear', align_corners=False)
+
+        # build outputs
+        outs = [
+            # self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
+            laterals[i] for i in range(used_backbone_levels)
+        ]
+        return tuple(outs)
+
+
+class PAN(FPN):
+    """Path Aggregation Network for Instance Segmentation.
+
+    This is an implementation of the `PAN in Path Aggregation Network
+    <https://arxiv.org/abs/1803.01534>`_.
+
+    Args:
+        in_channels (List[int]): Number of input channels per scale.
+        out_channels (int): Number of output channels (used at each scale)
+        num_outs (int): Number of output scales.
+        start_level (int): Index of the start input backbone level used to
+            build the feature pyramid. Default: 0.
+        end_level (int): Index of the end input backbone level (exclusive) to
+            build the feature pyramid. Default: -1, which means the last level.
+        add_extra_convs (bool): Whether to add conv layers on top of the
+            original feature maps. Default: False.
+        extra_convs_on_inputs (bool): Whether to apply extra conv on
+            the original feature from the backbone. Default: False.
+        relu_before_extra_convs (bool): Whether to apply relu before the extra
+            conv. Default: False.
+        no_norm_on_lateral (bool): Whether to apply norm on lateral.
+            Default: False.
+        conv_cfg (dict): Config dict for convolution layer. Default: None.
+        norm_cfg (dict): Config dict for normalization layer. Default: None.
+        act_cfg (str): Config dict for activation layer in ConvModule.
+            Default: None.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 num_outs,
+                 start_level=0,
+                 end_level=-1,
+                 conv_cfg=None,
+                 norm_cfg=None,
+                 activation=None):
+        super(PAN,
+              self).__init__(in_channels, out_channels, num_outs, start_level,
+                             end_level, conv_cfg, norm_cfg, activation)
+        self.init_weights()
+
+    def forward(self, inputs):
+        """Forward function."""
+        assert len(inputs) == len(self.in_channels)
+
+        # build laterals
+        laterals = [
+            lateral_conv(inputs[i + self.start_level])
+            for i, lateral_conv in enumerate(self.lateral_convs)
+        ]
+
+        # build top-down path
+        used_backbone_levels = len(laterals)
+        for i in range(used_backbone_levels - 1, 0, -1):
+            prev_shape = laterals[i - 1].shape[2:]
+            laterals[i - 1] += F.interpolate(
+                laterals[i], size=prev_shape, mode='bilinear', align_corners=False)
+
+        # build outputs
+        # part 1: from original levels
+        inter_outs = [
+            laterals[i] for i in range(used_backbone_levels)
+        ]
+
+        # part 2: add bottom-up path
+        for i in range(0, used_backbone_levels - 1):
+            prev_shape = inter_outs[i + 1].shape[2:]
+            inter_outs[i + 1] += F.interpolate(inter_outs[i], size=prev_shape, mode='bilinear', align_corners=False)
+
+        outs = []
+        outs.append(inter_outs[0])
+        outs.extend([
+            inter_outs[i] for i in range(1, used_backbone_levels)
+        ])
+        return tuple(outs)
+
+
+# if __name__ == '__main__':

models/module/init_weights.py

@@ -0,0 +1,44 @@
+"""
+from MMDetection
+"""
+import torch.nn as nn
+
+
+def kaiming_init(module,
+                 a=0,
+                 mode='fan_out',
+                 nonlinearity='relu',
+                 bias=0,
+                 distribution='normal'):
+    assert distribution in ['uniform', 'normal']
+    if distribution == 'uniform':
+        nn.init.kaiming_uniform_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    else:
+        nn.init.kaiming_normal_(
+            module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def xavier_init(module, gain=1, bias=0, distribution='normal'):
+    assert distribution in ['uniform', 'normal']
+    if distribution == 'uniform':
+        nn.init.xavier_uniform_(module.weight, gain=gain)
+    else:
+        nn.init.xavier_normal_(module.weight, gain=gain)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def normal_init(module, mean=0, std=1, bias=0):
+    nn.init.normal_(module.weight, mean, std)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)
+
+
+def constant_init(module, val, bias=0):
+    if hasattr(module, 'weight') and module.weight is not None:
+        nn.init.constant_(module.weight, val)
+    if hasattr(module, 'bias') and module.bias is not None:
+        nn.init.constant_(module.bias, bias)

models/module/norm.py

@@ -0,0 +1,55 @@
+import torch.nn as nn
+
+norm_cfg = {
+    # format: layer_type: (abbreviation, module)
+    'BN': ('bn', nn.BatchNorm2d),
+    'SyncBN': ('bn', nn.SyncBatchNorm),
+    'GN': ('gn', nn.GroupNorm),
+    # and potentially 'SN'
+}
+
+
+def build_norm_layer(cfg, num_features, postfix=''):
+    """ Build normalization layer
+
+    Args:
+        cfg (dict): cfg should contain:
+            type (str): identify norm layer type.
+            layer args: args needed to instantiate a norm layer.
+            requires_grad (bool): [optional] whether stop gradient updates
+        num_features (int): number of channels from input.
+        postfix (int, str): appended into norm abbreviation to
+            create named layer.
+
+    Returns:
+        name (str): abbreviation + postfix
+        layer (nn.Module): created norm layer
+    """
+    assert isinstance(cfg, dict) and 'type' in cfg
+    cfg_ = cfg.copy()
+
+    layer_type = cfg_.pop('type')
+    if layer_type not in norm_cfg:
+        raise KeyError('Unrecognized norm type {}'.format(layer_type))
+    else:
+        abbr, norm_layer = norm_cfg[layer_type]
+        if norm_layer is None:
+            raise NotImplementedError
+
+    assert isinstance(postfix, (int, str))
+    name = abbr + str(postfix)
+
+    requires_grad = cfg_.pop('requires_grad', True)
+    cfg_.setdefault('eps', 1e-5)
+    if layer_type != 'GN':
+        layer = norm_layer(num_features, **cfg_)
+        if layer_type == 'SyncBN':
+            layer._specify_ddp_gpu_num(1)
+    else:
+        assert 'num_groups' in cfg_
+        layer = norm_layer(num_channels=num_features, **cfg_)
+
+    for param in layer.parameters():
+        param.requires_grad = requires_grad
+
+    return name, layer

models/shufflenet2_att_m.py

@@ -0,0 +1,265 @@
+# -*- coding:utf-8 –*-
+import torch
+import torch.nn as nn
+from models.module.conv import DepthwiseConvModule
+
+
+class ShuffleV2Block(nn.Module):
+    def __init__(self, inp, oup, mid_channels, ksize, stride, attention='', ratio=16, loc='side', onnx=False):
+        super(ShuffleV2Block, self).__init__()
+        self.onnx = onnx
+        self.stride = stride
+        assert stride in [1, 2]
+
+        self.mid_channels = mid_channels
+        self.ksize = ksize
+        pad = ksize // 2
+        self.pad = pad
+        self.inp = inp
+
+        outputs = oup - inp
+
+        branch_main = [
+            # pw
+            nn.Conv2d(inp, mid_channels, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            # dw
+            nn.Conv2d(mid_channels, mid_channels, ksize, stride, pad, groups=mid_channels, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            # pw-linear
+            nn.Conv2d(mid_channels, outputs, 1, 1, 0, bias=False),
+            nn.BatchNorm2d(outputs),
+            nn.ReLU(inplace=True),
+        ]
+        self.branch_main = nn.Sequential(*branch_main)
+
+        if stride == 2:
+            branch_proj = [
+                # dw
+                nn.Conv2d(inp, inp, ksize, stride, pad, groups=inp, bias=False),
+                nn.BatchNorm2d(inp),
+                # pw-linear
+                nn.Conv2d(inp, inp, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(inp),
+                nn.ReLU(inplace=True),
+            ]
+            self.branch_proj = nn.Sequential(*branch_proj)
+        else:
+            self.branch_proj = None
+
+    def forward(self, old_x):
+        if self.stride == 1:
+            x_proj, x = self.channel_shuffle(old_x)
+        else:
+            x_proj = old_x
+            x_proj = self.branch_proj(x_proj)
+            x = old_x
+        x = self.branch_main(x)
+        x = torch.cat((x_proj, x), 1)
+        return x
+
+    def channel_shuffle(self, x):
+        batchsize, num_channels, height, width = x.data.size()
+        if self.onnx:
+            # 由于需要将onnx模型转换为ifx模型，ifx引擎以nchw(n=1)的格式存储数据，因此做shape变换时，尽量保证按nchw(n=1)来操作
+            x = x.reshape(1, batchsize * num_channels // 2, 2, height * width)
+            x = x.permute(0, 2, 1, 3)
+            z = num_channels // 2
+            x = x.reshape(1, -1, height, width)
+            # split时避免使用x[0]、x[1]的操作，尽量使用torch的算子来实现
+            x1, x2 = torch.split(x, split_size_or_sections=[z, z], dim=1)
+            return x1, x2
+        else:
+            x = x.reshape(batchsize * num_channels // 2, 2, height * width)
+            x = x.permute(1, 0, 2)
+            x = x.reshape(2, -1, num_channels // 2, height, width)
+            return x[0], x[1]
+
+
+class ShuffleNetV2(nn.Module):
+    def __init__(self, num_tasks=0, task_types=0, num_classes=[], out_channel=1024, model_size='0.5x', with_last_conv=True,
+                 attention='', loc='side', onnx=False, shuffle_block=None,
+                 stack_lite_head=0, lite_head_channels=-1):
+        super(ShuffleNetV2, self).__init__()
+        # print('model size is ', model_size)
+        assert len(num_classes) == num_tasks, f"num task must equal to length of classes list for every task"
+        
+        self.num_tasks = num_tasks
+        self.use_last_conv = with_last_conv
+        if isinstance(task_types, int):
+            task_types = [task_types] * num_tasks
+        if isinstance(stack_lite_head, int):
+            stack_lite_head = [stack_lite_head] * num_tasks
+        if isinstance(lite_head_channels, int):
+            lite_head_channels = [lite_head_channels] * num_tasks
+        self.task_types = task_types
+        self.stack_lite_head = stack_lite_head
+        self.lite_head_channels = lite_head_channels
+        self.onnx = onnx
+        self.stage_repeats = [4, 8, 4]
+        self.model_size = model_size
+        if model_size == '0.5x':
+            self.stage_out_channels = [-1, 24, 48, 96, 192] + [out_channel]
+        elif model_size == '1.0x':
+            self.stage_out_channels = [-1, 24, 116, 232, 464] + [out_channel]
+        elif model_size == '1.5x':
+            self.stage_out_channels = [-1, 24, 176, 352, 704] + [out_channel]
+        elif model_size == '2.0x':
+            self.stage_out_channels = [-1, 24, 244, 488, 976] + [out_channel]
+        else:
+            raise NotImplementedError
+
+        shuffle_block = ShuffleV2Block if shuffle_block is None else shuffle_block
+        # building first layer
+        input_channel = self.stage_out_channels[1]
+        self.first_conv = nn.Sequential(
+            nn.Conv2d(3, input_channel, 3, 2, 1, bias=False),
+            nn.BatchNorm2d(input_channel),
+            nn.ReLU(inplace=True),
+        )
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.features = []
+        for idxstage in range(len(self.stage_repeats)):
+            numrepeat = self.stage_repeats[idxstage]
+            output_channel = self.stage_out_channels[idxstage + 2]
+
+            for i in range(numrepeat):
+                if i == 0:
+                    self.features.append(
+                        shuffle_block(input_channel, output_channel, mid_channels=output_channel // 2, ksize=3,
+                                      stride=2, attention=attention, ratio=16, loc=loc, onnx=onnx))
+                else:
+                    self.features.append(
+                        shuffle_block(input_channel//2, output_channel, mid_channels=output_channel//2, ksize=3,
+                                      stride=1, attention=attention, ratio=16, loc=loc, onnx=onnx))
+
+                input_channel = output_channel
+
+        self.features = nn.Sequential(*self.features)
+
+        self.classifier_inchannels = self.stage_out_channels[-2]
+        if with_last_conv:
+            self.classifier_inchannels = self.stage_out_channels[-1]
+            self.conv_last = nn.Sequential(
+                nn.Conv2d(input_channel, self.classifier_inchannels, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(self.stage_out_channels[-1]),
+                nn.ReLU(inplace=True))
+
+        self.globalpool = nn.AdaptiveAvgPool2d(output_size=1)
+
+        self.lite_head_channels = [self.classifier_inchannels if v == -1 else v for v in self.lite_head_channels]
+        for ti in range(self.num_tasks):
+            if self.stack_lite_head[ti]:
+                lite_head = []
+                for j in range(self.stack_lite_head[ti]):
+                    ins = self.classifier_inchannels if j == 0 else self.lite_head_channels[ti]
+                    outs = self.classifier_inchannels if j == self.stack_lite_head[ti]-1 else self.lite_head_channels[ti]
+                    lite_head.append(DepthwiseConvModule(ins, outs, 3, 1, 1))
+                lite_head = nn.Sequential(*lite_head)
+                self.add_module(f"lite_head{ti}", lite_head)
+            classifier = nn.Sequential(nn.Linear(self.classifier_inchannels, num_classes[ti], bias=False))
+            self.add_module(f"classifier{ti}", classifier)
+
+        # self.loss_weights = nn.Parameter(torch.ones(num_tasks), requires_grad=True)
+
+        self._initialize_weights()
+
+    def _forward_impl(self, x):
+        x = self.first_conv(x)
+        x = self.maxpool(x)
+        x = self.features(x)
+
+        if self.use_last_conv:
+            x = self.conv_last(x)
+
+        output = []
+        for ti in range(self.num_tasks):
+            if self.stack_lite_head[ti]:
+                c_x = getattr(self, f"lite_head{ti}")(x)
+                c_x = self.globalpool(c_x)
+                c_x = c_x.contiguous().view(-1, self.classifier_inchannels)
+                c_x = getattr(self, f"classifier{ti}")(c_x)
+            else:
+                c_x = self.globalpool(x)
+                c_x = c_x.contiguous().view(-1, self.classifier_inchannels)
+                c_x = getattr(self, f"classifier{ti}")(c_x)
+            
+            if self.onnx:
+                if self.task_types[ti] == 0:
+                    c_x = torch.softmax(c_x, dim=1)
+                elif self.task_types[ti] == 1:
+                    c_x *= (0.05/3)
+                elif self.task_types[ti] == 2:
+                    c_x *= (0.7/3)
+                elif self.task_types[ti] == 3:
+                    c_x *= (0.1/2)
+                else:
+                    raise NotImplementedError(f"task_type only support [0, 1, 2, 3], current {self.task_types[ti]}")
+            
+            output.append(c_x)
+        return output
+
+    def forward(self, x):
+        output = self._forward_impl(x)
+        return output
+
+    def _initialize_weights(self):
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_in')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_in')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def get_last_share_layer(self):
+        assert self.use_last_conv, "Current implement need 'with_last_conv=True'"
+        return self.conv_last[0]
+
+
+if __name__ == '__main__':
+    from models.create_model import speed, model_info
+    from torchstat import stat
+
+    # a = se_resnet101()
+    # b = shufflenet_v2_x0_5()
+    input_size = (160, 160)
+    shufflenet2 = ShuffleNetV2(out_channel=192,
+        num_tasks=10, task_types=1, num_classes=[5, 3, 3, 3, 3, 1, 1, 1, 1, 1], model_size='0.5x', with_last_conv=0,
+        stack_lite_head=1, lite_head_channels=-1, onnx=True)
+    #
+    model_info(shufflenet2, img_size=input_size, verbose=False)
+    speed(shufflenet2, 'shufflenet2', size=input_size, device_type='cpu')
+    stat(shufflenet2, input_size=(3, 160, 160))
+
+
+
+    # shufflenet2.eval()
+    #
+    # # example = torch.randn(1, 3, input_size[1], input_size[0])
+    # # torch.onnx.export(
+    # #     shufflenet2,  # model being run
+    # #     example,  # model input (or a tuple for multiple inputs)
+    #     '1.onnx',
+    #     verbose=False,
+    #     # store the trained parameter weights inside the model file
+    #     training=False,
+    #     input_names=['input'],
+    #     output_names=['output'],
+    #     do_constant_folding=True
+    # )

utils/common.py

@@ -0,0 +1,227 @@
+# -*-coding:utf-8-*-
+import colorsys
+import datetime
+import functools
+import signal
+import sys
+from contextlib import contextmanager
+
+import numpy as np
+
+import time
+
+
+# Generate colors
+def gen_color(n):
+    hsv_tuples = [(x / n, 1., 1.)
+                  for x in range(n)]
+    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
+    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
+    # np.random.seed(10101)  # Fixed seed for consistent colors across runs.
+    np.random.shuffle(colors)  # Shuffle colors to decorrelate adjacent classes.
+    # np.random.seed(None)  # Reset seed to default.
+    return colors
+
+
+# pylint: disable=W0232
+class Color:
+    GRAY = 30
+    RED = 31
+    GREEN = 32
+    YELLOW = 33
+    BLUE = 34
+    MAGENTA = 35
+    CYAN = 36
+    WHITE = 67
+    CRIMSON = 38
+
+
+# 返回字符串的输出格式码，调整前景色背景色、加粗等
+def colorize(num, string, bold=False, highlight=False):
+    assert isinstance(num, int)
+    attr = []
+    if bold:
+        attr.append('1')
+    if highlight and num == 67:
+        num += 30
+    if highlight and num != 67:
+        num += 60
+    attr.append(str(num))
+    # \x1b[显示方式;前景色;背景色m +"输出内容"+\x1b[0m
+    # ; 的顺序可以改变
+    return '\x1b[%sm%s\x1b[0m' % (';'.join(attr), string)
+
+
+def colorprint(colorcode, text, o=sys.stdout, bold=False, highlight=False, end='\n'):
+    o.write(colorize(colorcode, text, bold=bold, highlight=highlight) + end)
+
+
+def cprint(text, colorcode=67, bold=False, highlight=False, end='\n', prefix=None,
+           pre_color=34, pre_bold=True, pre_high=True, pre_end=': '):
+    prefix = str(prefix).rstrip() if prefix is not None else prefix
+    if prefix is not None:
+        prefix = prefix.rstrip(':') if ':' in pre_end else prefix
+        prefix += pre_end
+        colorprint(pre_color, prefix, sys.stdout, pre_bold, pre_high, end='')
+    colorprint(colorcode, text, bold=bold, highlight=highlight, end=end)
+
+
+def log_warn(msg):
+    cprint(msg, colorcode=33, prefix='Warning', pre_color=33, highlight=True)
+
+
+def log_error(msg):
+    cprint(msg, colorcode=31, prefix='Error', pre_color=31, highlight=True)
+
+
+def now_time():
+    return datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+
+
+# http://stackoverflow.com/questions/366682/how-to-limit-execution-time-of-a-function-call-in-python
+class TimeoutException(Exception):
+    def __init__(self, msg):
+        self.msg = msg
+
+
+# 上下文管理器，装饰函数func, func()中必须有yield.使用方法：
+# with func as f:
+#   do()
+# with语句执行yield之前的代码，然后执行f和do，最后执行yield之后的代码
+@contextmanager
+def time_limit(seconds):
+    # 这是一个限制程序运行时间的函数，如果超出预设时间，会报错
+    def signal_handler(signum, frame):  # 信号处理函数必须有这两个参数
+        raise TimeoutException(colorize(Color.RED, "Timed out! Retry again ...", highlight=True))
+    signal.signal(signal.SIGALRM, signal_handler)  # 接收信号，signal.SIGALRM 为信号，signal_handler处理信号的函数
+    signal.alarm(seconds)  # 如果seconds是非0，这个函数则响应一个SIGALRM信号并在seconds秒后发送到该进程。
+    try:
+        yield
+    finally:
+        signal.alarm(0)  # 中断信号
+
+
+def clock(func):  # 加强版
+    @functools.wraps(func)  # 把func的属性复制给clocked，避免func的name变成clocked
+    def clocked(*args, **kwargs):  # 内部函数可接受任意个参数，以及支持关键字参数
+        t0 = time.perf_counter()
+        result = func(*args, **kwargs)  # clocked闭包中包含自由变量func
+        elapsed = time.perf_counter() - t0
+        name = func.__name__
+        arg_lst = []
+        if args:
+            arg_lst.append(', '.join(repr(arg) for arg in args))
+        if kwargs:
+            pairs = ['%s=%r' % (k, w) for k, w in sorted(kwargs.items())]
+            arg_lst.append(', '.join(pairs))
+        arg_str = ', '.join(arg_lst)
+        colorprint(Color.GREEN, '[%0.8fs] %s(%s) -> %r' % (elapsed, name, arg_str, result))
+        return result
+    return clocked  # 返回clocked，取代被装饰的函数
+
+
+DEFAULT_FMT = '[{elapsed:0.8f}s] {name}({args}) -> {result}'
+
+
+def clock_custom(fmt=DEFAULT_FMT, color=Color.GREEN):
+    def decorate(func):  # 加强版
+        @functools.wraps(func)  # 把func的属性复制给clocked，避免func的name变成clocked
+        def clocked(*_args, **_kwargs):  # 内部函数可接受任意个参数，以及支持关键字参数
+            t0 = time.perf_counter()
+            _result = func(*_args, **_kwargs)  # clocked闭包中包含自由变量func
+            elapsed = time.perf_counter() - t0
+            name = func.__name__
+            _arg_lst = []
+            if _args:
+                _arg_lst.append(', '.join(repr(arg) for arg in _args))
+            if _kwargs:
+                _pairs = ['%s=%r' % (k, w) for k, w in sorted(_kwargs.items())]
+                _arg_lst.append(', '.join(_pairs))
+            args = ', '.join(_arg_lst)
+            result = repr(_result)  # 字符串形式
+            colorprint(color, fmt.format(**locals()))  # 使用clocked的局部变量
+            return _result  # 返回原函数的结果
+        return clocked  # 返回clocked，取代被装饰的函数
+    return decorate  # 装饰器工厂必须返回装饰器
+
+
+def colorstr(*input):
+    # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
+    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
+    colors = {'black': '\033[30m',  # basic colors
+              'red': '\033[31m',
+              'green': '\033[32m',
+              'yellow': '\033[33m',
+              'blue': '\033[34m',
+              'magenta': '\033[35m',
+              'cyan': '\033[36m',
+              'white': '\033[37m',
+              'bright_black': '\033[90m',  # bright colors
+              'bright_red': '\033[91m',
+              'bright_green': '\033[92m',
+              'bright_yellow': '\033[93m',
+              'bright_blue': '\033[94m',
+              'bright_magenta': '\033[95m',
+              'bright_cyan': '\033[96m',
+              'bright_white': '\033[97m',
+              'end': '\033[0m',  # misc
+              'bold': '\033[1m',
+              'underline': '\033[4m'}
+    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
+
+
+class Logger:
+    def __init__(self, file_path):
+        self.log_file_path = file_path
+        self.log_file = None
+        self.color = {
+            'R': Color.RED,
+            'B': Color.BLUE,
+            'G': Color.GREEN,
+            'Y': Color.YELLOW
+        }
+
+    def start(self):
+        self.log_file = open(self.log_file_path, 'w', encoding='utf-8')
+
+    def close(self):
+        if self.log_file is not None:
+            self.log_file.close()
+        return
+
+    def info(self, text, color='W', prefix=None, pre_color='B', pre_end=': ', prints=True):
+        assert self.log_file is not None, "Please firstly confirm 'logger.start()' method"
+        color_code = self.color.get(color, Color.WHITE)
+        pre_color = self.color.get(pre_color, Color.BLUE)
+        prefix = str(prefix).rstrip() if prefix is not None else prefix
+        if prefix is not None:
+            prefix = prefix.rstrip(':') if ':' in pre_end else prefix
+            prefix += pre_end
+        self.log_file.write(f"{prefix if prefix is not None else ''}{text}\n")
+        if prints:
+            cprint(text, color_code, prefix=prefix, pre_color=pre_color)
+
+    def error(self, text):
+        assert self.log_file is not None, "Please firstly confirm 'logger.start()' method"
+        # self.log_file.write(text + '\n')
+        log_error(text)
+
+    def warn(self, text):
+        assert self.log_file is not None, "Please firstly confirm 'logger.start()' method"
+        # self.log_file.write(text + '\n')
+        log_warn(text)
+
+
+def r(val):
+    return int(np.random.random() * val)
+
+
+# if __name__ == '__main__':
+#     #print(colorize(31, 'I am fine!', bold=True, highlight=True))
+#     #colorprint(35, 'I am fine!')
+#     #error('get out')
+#     import time
+#
+#     # ends after 5 seconds
+#     with time_limit(2):
+#         time.sleep(3)

utils/export_util.py

@@ -0,0 +1,118 @@
+# -*- coding:utf-8 –*-
+# tensorflow=2.5.0 onnx-tf=1.9.0
+import os
+import time
+import onnxruntime
+# import tensorflow as tf
+
+
+def onnx_export_pipeline(onnx_path, export_func, simplify=True, onnx_sim_path=None, pack=True, view_net=True):
+    if not os.path.exists(os.path.dirname(onnx_path)):
+        os.makedirs(os.path.dirname(onnx_path))
+
+    # 导出
+    export_func(export_name=onnx_path)  # 转化onnx的函数，与各自的配置密切相关
+    if not simplify:
+        print(f"Onnx model export in '{onnx_path}'")
+
+    if simplify:  # 算子简化
+        time.sleep(1)
+        onnx_sim_path = onnx_sim_path if onnx_sim_path else onnx_path
+        os.system(f"python -m onnxsim {onnx_path} {onnx_sim_path}")
+        print(f"Simplify onnx model export in '{onnx_sim_path}'")
+
+    if pack:      # 压缩打包
+        time.sleep(1)
+        src_path = onnx_sim_path if simplify else onnx_path
+        src_dir, src_name = os.path.dirname(src_path), os.path.basename(src_path)
+        os.system(f'cd {src_dir} && tar -zcf {src_name}.tgz {src_name}')
+        print(f"TGZ file save in '{src_path}.tgz'")
+
+    if view_net:  # 查看网络结构
+        import netron
+        src_path = onnx_sim_path if simplify else onnx_path
+        netron.start(src_path)
+
+
+class ONNXModel:
+    def __init__(self, onnx_path):
+        """
+        :param onnx_path:
+        """
+        self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+        self.input_name = self.get_input_name(self.onnx_session)
+        self.output_name = self.get_output_name(self.onnx_session)
+        print(f"loading {onnx_path}")
+
+    @staticmethod
+    def get_output_name(onnx_session):
+        """
+        output_name = onnx_session.get_outputs()[0].name
+        :param onnx_session:
+        :return:
+        """
+        output_name = []
+        for node in onnx_session.get_outputs():
+            output_name.append(node.name)
+        return output_name
+
+    @staticmethod
+    def get_input_name(onnx_session):
+        """
+        input_name = onnx_session.get_inputs()[0].name
+        :param onnx_session:
+        :return:
+        """
+        input_name = []
+        for node in onnx_session.get_inputs():
+            input_name.append(node.name)
+        return input_name
+
+    @staticmethod
+    def get_input_feed(input_name, image_numpy):
+        """
+        input_feed={self.input_name: image_numpy}
+        :param input_name:
+        :param image_numpy:
+        :return:
+        """
+        input_feed = {}
+        for name in input_name:
+            input_feed[name] = image_numpy
+        return input_feed
+
+    def forward(self, image_numpy):
+        '''
+        # image_numpy = image.transpose(2, 0, 1)
+        # image_numpy = image_numpy[np.newaxis, :]
+        # onnx_session.run([output_name], {input_name: x})
+        # :param image_numpy:
+        # :return:
+        '''
+        input_feed = self.get_input_feed(self.input_name, image_numpy)
+        outputs = self.onnx_session.run(self.output_name, input_feed=input_feed)
+        return outputs
+
+
+def to_numpy(tensor):
+    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor
+
+
+def export_tflite(onnx_path, tflite_path):
+    save_model_dir = os.path.splitext(tflite_path)[0]
+    onnx_to_pb_cmd = f'onnx-tf convert -i {onnx_path} -o {save_model_dir}'
+    os.system(onnx_to_pb_cmd)
+    print(f"Convert onnx model to pb in '{save_model_dir}'")
+
+    time.sleep(1)
+    # make a converter object from the saved tensorflow file
+    converter = tf.lite.TFLiteConverter.from_saved_model(save_model_dir)
+    # tell converter which type of optimization techniques to use
+    # to view the best option for optimization read documentation of tflite about optimization
+    # go to this link https://www.tensorflow.org/lite/guide/get_started#4_optimize_your_model_optional
+    converter.optimizations = [tf.lite.Optimize.DEFAULT]
+    # convert the model
+    tf_lite_model = converter.convert()
+    # save the converted model
+    open(tflite_path, 'wb').write(tf_lite_model)
+    print(f"TFlite model save in '{tflite_path}'")

utils/images.py

@@ -0,0 +1,539 @@
+# -*- coding:utf-8 –*-
+import random
+from math import *
+
+import cv2
+from PIL import ImageDraw, Image, ImageFont
+
+from utils.common import *
+from utils.multiprogress import MultiThreading
+
+
+def show(image, name="Press 'q' exit"):
+    if not hasattr(image, 'shape'):
+        img_array = cv2.cvtColor(np.asarray(image), cv2.COLOR_RGB2BGR)
+    else:
+        img_array = image.copy()
+    cv2.imshow(name, img_array)
+    cv2.moveWindow(name, 0, 0)
+    key = cv2.waitKey(0) & 0xEFFFFF
+    if key == ord('q'):
+        cv2.destroyWindow(name)
+
+
+def crop_face_square(img, box, pad=5):
+    img_h, img_w, img_c = img.shape
+    if box != [0, 0, 0, 0]:
+        x0, y0, x1, y1 = box
+        w, h = x1-x0, y1-y0
+        # 扩充为正方形
+        long_side = min(max(h, w) + int(2*pad), img_h)
+        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
+        crop_x0, crop_y0 = max(x0 - w_add, 0), max(y0 - h_add, 0)
+        crop_x1, crop_y1 = min(crop_x0 + long_side, img_w), min(crop_y0 + long_side, img_h)
+        crop_x0, crop_y0 = crop_x1 - long_side, crop_y1 - long_side
+        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]
+    else:
+        # print('No detected box, crop right rect.')
+        if img_h == 960:
+            img = img[60:780, :]
+            img_h, img_w = img.shape[:2]
+        return img[:, img_w - img_h:img_w], [0, 0, 0, 0]
+
+
+def crop_face_square_rate(img, box, rate=0.1):
+    img_h, img_w, img_c = img.shape
+    if box != [0, 0, 0, 0]:
+        x0, y0, x1, y1 = box
+        w, h = x1-x0, y1-y0
+        # 扩充为正方形
+        pad = max(w, h) * rate
+        long_side = min(max(h, w) + int(2*pad), img_h)
+        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
+        crop_x0, crop_y0 = max(x0 - w_add, 0), max(y0 - h_add, 0)
+        crop_x1, crop_y1 = min(crop_x0 + long_side, img_w), min(crop_y0 + long_side, img_h)
+        crop_x0, crop_y0 = crop_x1 - long_side, crop_y1 - long_side
+        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]
+    else:
+        if img_h == 960:
+            crop_x0, crop_x1 = img_w - 720, img_w
+            crop_y0, crop_y1 = 60, 780
+        else:
+            crop_x0, crop_x1 = img_w - img_h, img_w
+            crop_y0, crop_y1 = 0, img_h
+        return img[crop_y0:crop_y1, crop_x0:crop_x1], [crop_x0, crop_y0, crop_x1, crop_y1]
+
+
+def expand_box_rate(img, box, rate=0.1):
+    img_h, img_w, img_c = img.shape
+    if box != [0, 0, 0, 0]:
+        x0, y0, x1, y1 = box
+        w, h = x1-x0, y1-y0
+        # 扩充为正方形
+        pad = max(w, h) * rate
+        long_side = min(max(h, w) + int(2*pad), img_h)
+        w_add, h_add = (long_side - w) // 2, (long_side - h) // 2
+        crop_x0, crop_y0 = x0 - w_add, y0 - h_add
+        crop_x1, crop_y1 = crop_x0 + long_side, crop_y0 + long_side
+        return [crop_x0, crop_y0, crop_x1, crop_y1]
+    else:
+        if img_h == 960:
+            crop_x0, crop_x1 = img_w - 720, img_w
+            crop_y0, crop_y1 = 60, 780
+        else:
+            crop_x0, crop_x1 = img_w - img_h, img_w
+            crop_y0, crop_y1 = 0, img_h
+        return [crop_x0, crop_y0, crop_x1, crop_y1]
+
+
+def crop_with_pad(img, crop_box):
+    img_h, img_w, img_c = img.shape
+    x0, y0, x1, y1 = crop_box
+    w, h = x1 - x0, y1 - y0
+    if tuple(crop_box) == (0, 0, 0, 0) or w <= 50 or h <= 50:  # 背景，裁右半图
+        if img_h == 960:
+            crop_x0, crop_x1 = img_w - 720, img_w
+            crop_y0, crop_y1 = 60, 780
+        else:
+            crop_x0, crop_x1 = img_w - img_h, img_w
+            crop_y0, crop_y1 = 0, img_h
+        crop_img = img[crop_y0:crop_y1, crop_x0:crop_x1]
+        return crop_img
+    else:
+        crop_x0, crop_y0 = max(x0, 0), max(y0, 0)
+        crop_x1, crop_y1 = min(x1, img_w), min(y1, img_h)
+        left, top, right, bottom = crop_x0 - x0, crop_y0 - y0, x1 - crop_x1, y1 - crop_y1
+        crop_img = img[crop_y0:crop_y1, crop_x0:crop_x1]
+        crop_img = cv2.copyMakeBorder(crop_img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(128, 128, 128))
+        return crop_img
+
+
+def clip_paste_rect(bg_img, fg_img, loc, refine=True):
+    """
+    选取粘贴区域，不要超出背景范围
+    loc: 中心点(cx, cy)
+    """
+    bg_h, bg_w = bg_img.shape[:2]
+    fg_h, fg_w = fg_img.shape[:2]
+    fg_h = fg_h - 1 if fg_h % 2 else fg_h
+    fg_w = fg_w - 1 if fg_w % 2 else fg_w
+    cx, cy = loc
+    left, top, right, bottom = cx - fg_w // 2, cy - fg_h // 2, cx + fg_w // 2, cy + fg_h // 2
+
+    if refine:
+        right, bottom = min(right, bg_w), min(bottom, bg_h)
+        left, top = right - fg_w, bottom - fg_h
+        left, top = max(0, left), max(0, top)
+        right, bottom = left + fg_w, top + fg_h
+
+    plot_x1, plot_y1, plot_x2, plot_y2 = left, top, right, bottom
+    use_x1, use_y1, use_x2, use_y2 = 0, 0, fg_w, fg_h
+
+    if left < 0:
+        plot_x1, use_x1 = 0, -left
+    if top < 0:
+        plot_y1, use_y1 = 0, -top
+    if right > bg_w:
+        plot_x2, use_x2 = bg_w, fg_w - (right - bg_w)
+    if bottom > bg_h:
+        plot_y2, use_y2 = bg_h, fg_h - (bottom - bg_h)
+
+    use_bg = bg_img[plot_y1:plot_y2, plot_x1:plot_x2]
+    use_fg = fg_img[use_y1:use_y2, use_x1:use_x2]
+    window = (plot_x1, plot_y1, plot_x2, plot_y2)
+
+    return use_bg, use_fg, window
+
+
+# @clock_custom(fmt=DEFAULT_FMT)
+def paste(bg_img, fg_img, loc, trans_thresh=1, refine=True):
+    """
+    贴图
+    loc: center (cx, cy)
+    """
+    use_bg, use_fg, window = clip_paste_rect(bg_img, fg_img, loc, refine)
+    plot_x1, plot_y1, plot_x2, plot_y2 = window
+    b, g, r, a = cv2.split(use_fg)
+    a[a > 0] = 255
+    a = np.dstack([a, a, a]) * trans_thresh
+    use_bg = use_bg * (255.0 - a) / 255
+    use_bg += use_fg[:, :, :3] * (a / 255)
+    bg_img[plot_y1:plot_y2, plot_x1:plot_x2] = use_bg.astype('uint8')
+    return bg_img
+
+
+def put_chinese_text(img, text, position, font, text_color=(0, 255, 0)):
+    if isinstance(img, np.ndarray):  # 判断是否OpenCV图片类型
+        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    # 创建一个可以在给定图像上绘图的对象
+    draw = ImageDraw.Draw(img)
+    draw.text(position, text, text_color, font=font)
+    # 转换回OpenCV格式
+    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
+
+
+def compute_mean_std(img_list, preprocess=None, workers=12):
+    """
+    计算公式：
+    S^2 = sum((x - x')^2) / N = sum(x^2 + x'^2 - 2xx') / N
+        = (sum(x^2) + sum(x'^2) - 2x'*sum(x)) / N
+        = (sum(x^2) + N*(x'^2) - 2x'*(N * x')) / N
+        = (sum(x^2) - N * (x'^2)) / N
+        = sum(x^2) / N - (sum(x) / N)^2
+        = mean(x^2) - mean(x)^2 = E(x^2) - E(x)^2
+    :param img_list:
+    :param workers:
+    :param preprocess: 图像预处理函数，需要有BGR->RGB以及归一化操作
+    :return: RGB means, stds
+    """
+
+    def cal_func(info):
+        i, img_path = info
+        img = cv2.imread(img_path)
+        if preprocess:
+            img = preprocess(img)
+        else:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = img.astype(np.float32) / 255
+        img = img.reshape((-1, 3))
+        channel_mean = np.mean(img, axis=0)
+        channel_var = np.mean(img**2, axis=0)
+        if i % 1000 == 0:
+            print(f"{i}/{len(img_list)}")
+        return channel_mean, channel_var
+
+    exe = MultiThreading(list(enumerate(img_list)), workers=min(workers, len(img_list)))
+    res = exe.run(cal_func)
+    all_means = np.array([r[0] for r in res])
+    all_vars = np.array([r[1] for r in res])
+    means = np.mean(all_means, axis=0)
+    vars = np.mean(all_vars, axis=0)
+    stds = np.sqrt(vars - means ** 2)
+
+    return means, stds
+
+
+def draw_province(f, val):
+    img = Image.new("RGB", (45, 70), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    draw.text((0, 3), val, (0, 0, 0), font=f)
+    img = img.resize((23, 70))
+    char = np.array(img)
+    return char
+
+
+def draw_chars(f, val):
+    img = Image.new("RGB", (23, 70), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    draw.text((0, 2), val, (0, 0, 0), font=f)
+    A = np.array(img)
+    return A
+
+
+# 双层车牌
+def draw_province_double(f, val):
+    img = Image.new("RGB", (60, 60), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    draw.text((0, -12), val, (0, 0, 0), font=f)
+    img = img.resize((80, 60))
+    # img.show()
+    char = np.array(img)
+    return char
+
+
+def draw_chars_ceil(f, val):
+    img = Image.new("RGB", (30, 45), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    draw.text((1, -12), val, (0, 0, 0), font=f)
+    img = img.resize((80, 60))
+    # img.show()
+    char = np.array(img)
+    return char
+
+
+def draw_chars_floor(f, val):
+    img = Image.new("RGB", (30, 45), (255, 255, 255))
+    draw = ImageDraw.Draw(img)
+    draw.text((1, -12), val, (0, 0, 0), font=f)
+    img = img.resize((65, 110))
+    # img.show()
+    char = np.array(img)
+    return char
+
+
+# 图片做旧处理
+def add_smudginess(img):
+    smu = cv2.imread('data/bgp/smu.jpg')
+    img_h, img_w = img.shape[:2]
+    rows = r(smu.shape[0] - img_h)
+    cols = r(smu.shape[1] - img_w)
+    adder = smu[rows:rows + img_h, cols:cols + img_w]
+    adder = cv2.resize(adder, (img_w, img_h))
+    adder = cv2.bitwise_not(adder)
+    val = random.random() * 0.5
+    img = cv2.addWeighted(img, 1 - val, adder, val, 0.0)
+    return img
+
+
+def rot(img, angel, shape, max_angel):
+    """ 使图像轻微的畸变
+        img 输入图像
+        factor 畸变的参数
+        size 为图片的目标尺寸
+    """
+    size_o = [shape[1], shape[0]]
+    size = (shape[1] + int(shape[0] * sin((float(max_angel)/180) * 3.14)), shape[0])
+    interval = abs(int(sin((float(angel) / 180) * 3.14) * shape[0]))
+    pts1 = np.float32([[0, 0], [0, size_o[1]], [size_o[0], 0], [size_o[0], size_o[1]]])
+    if angel > 0:
+        pts2 = np.float32([[interval, 0], [0, size[1]], [size[0], 0], [size[0]-interval, size_o[1]]])
+    else:
+        pts2 = np.float32([[0, 0], [interval, size[1]], [size[0]-interval, 0], [size[0], size_o[1]]])
+    M = cv2.getPerspectiveTransform(pts1, pts2)
+    dst = cv2.warpPerspective(img, M, size)
+    return dst
+
+
+def random_rot(img, factor, size):
+    shape = size
+    pts1 = np.float32([[0, 0], [0, shape[0]], [shape[1], 0], [shape[1], shape[0]]])
+    pts2 = np.float32([[r(factor), r(factor)],
+                        [r(factor), shape[0] - r(factor)],
+                        [shape[1] - r(factor), r(factor)],
+                        [shape[1] - r(factor), shape[0] - r(factor)]])
+    M = cv2.getPerspectiveTransform(pts1, pts2)
+    dst = cv2.warpPerspective(img, M, size)
+    return dst
+
+
+def random_rot_expend(img, factor, size):
+    height, width = size
+    degree = factor
+    # 旋转后的尺寸
+    n_h = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))
+    n_w = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))
+    M = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)
+    M[0, 2] += (n_w - width) / 2  # 重点在这步，目前不懂为什么加这步
+    M[1, 2] += (n_h - height) / 2  # 重点在这步
+    rot_img = cv2.warpAffine(img, M, (n_w, n_h), borderValue=(0, 0, 0))
+    return rot_img
+
+
+def random_rot_keep_size_left_right(img, factor1, factor2, size):
+    # 透视变换 factor1 大于0 factor2 可正可负
+    shape = size
+    width = shape[0]
+    height = shape[1]
+    pts1 = np.float32([[0, 0], [width - 1, 0], [0, height-1], [width - 1, height-1]])
+    point1_x = 0
+    point1_y = 0
+    point2_x = width - r(factor1)
+    point2_y = r(factor2)  # shape[0] - r(factor)
+    point3_x = 0  # shape[1] - r(factor)
+    point3_y = height  # r(factor)
+    point4_x = width - r(factor1)  # shape[1]
+    point4_y = height - r(factor2)  # shape[0]
+    max_x = max(point2_x, point4_x)
+    max_y = point3_y
+    if factor2 < 0:
+        point1_x = 0
+        point1_y = 0 - r(factor2)
+        point2_x = width - r(factor1)
+        point2_y = 0
+        point3_x = 0
+        point3_y = height + r(factor2)
+        point4_x = width - r(factor1)
+        point4_y = height
+        max_x = max(point2_x, point4_x)
+        max_y = point4_y
+    pts2 = np.float32([[point1_x, point1_y],
+                       [point2_x, point2_y],
+                       [point3_x, point3_y],
+                       [point4_x, point4_y]])
+    M = cv2.getPerspectiveTransform(pts1, pts2)
+    size2 = (max_x, max_y)
+    dst = cv2.warpPerspective(img, M, size2)  # cv2.warpPerspective(img, M, size)
+    return dst
+
+
+# 仿射变换
+def affine_transform(img, factor, size):
+    shape = size
+    pts1 = np.float32([[0, shape[0]], [shape[1], 0], [shape[1], shape[0]]])
+    pts2 = np.float32([[r(factor), shape[0] - r(factor)],
+                       [shape[1] - r(factor), r(factor)],
+                       [shape[1], shape[0]]])  # [shape[1] - r(factor), shape[0] - r(factor)]])
+    M = cv2.getAffineTransform(pts1, pts2)
+    dst = cv2.warpAffine(img, M, size)
+    return dst
+
+
+# 腐蚀
+def cv_erode(img, factor):
+    value = r(factor)+1
+    kernel = np.ones((value, value), np.uint8) * factor
+    erosion = cv2.erode(img, kernel, iterations=1)
+    return erosion
+
+
+# 膨胀
+def cv_dilate(img, factor):
+    value = r(factor)+1
+    kernel = np.ones((value, value), np.uint8)
+    dilate = cv2.dilate(img, kernel, iterations=1)
+    return dilate
+
+
+def add_random_noise(img, factor):
+    value = r(factor)
+    for k in range(value):  # Create 5000 noisy pixels
+        i = random.randint(0, img.shape[0] - 1)
+        j = random.randint(0, img.shape[1] - 1)
+        color = (random.randrange(256), random.randrange(256), random.randrange(256))
+        img[i, j] = color
+    return img
+
+
+# 生成卷积核和锚点
+def gen_kernel_anchor(length, angle):
+    half = length / 2
+    EPS = np.finfo(float).eps
+    alpha = (angle - floor(angle / 180) * 180) / 180 * pi
+    cosalpha = cos(alpha)
+    sinalpha = sin(alpha)
+    if cosalpha < 0:
+        xsign = -1
+    elif angle == 90:
+        xsign = 0
+    else:
+        xsign = 1
+    psfwdt = 1
+    # 模糊核大小
+    sx = int(fabs(length * cosalpha + psfwdt * xsign - length * EPS))
+    sy = int(fabs(length * sinalpha + psfwdt - length * EPS))
+    psf1 = np.zeros((sy, sx))
+    # psf1是左上角的权值较大，越往右下角权值越小的核。
+    # 这时运动像是从右下角到左上角移动
+    for i in range(0, sy):
+        for j in range(0, sx):
+            psf1[i][j] = i * fabs(cosalpha) - j * sinalpha
+            rad = sqrt(i * i + j * j)
+            if rad >= half and fabs(psf1[i][j]) <= psfwdt:
+                temp = half - fabs((j + psf1[i][j] * sinalpha) / cosalpha)
+                psf1[i][j] = sqrt(psf1[i][j] * psf1[i][j] + temp * temp)
+            psf1[i][j] = psfwdt + EPS - fabs(psf1[i][j])
+            if psf1[i][j] < 0:
+                psf1[i][j] = 0
+    # 运动方向是往左上运动，锚点在（0，0）
+    anchor = (0, 0)
+    # 运动方向是往右上角移动，锚点一个在右上角
+    # 同时，左右翻转核函数，使得越靠近锚点，权值越大
+    if 0 < angle < 90:
+        psf1 = np.fliplr(psf1)
+        anchor = (psf1.shape[1] - 1, 0)
+    elif -90 < angle < 0:  # 同理：往右下角移动
+        psf1 = np.flipud(psf1)
+        psf1 = np.fliplr(psf1)
+        anchor = (psf1.shape[1] - 1, psf1.shape[0] - 1)
+    elif angle < -90:  # 同理：往左下角移动
+        psf1 = np.flipud(psf1)
+        anchor = (0, psf1.shape[0] - 1)
+    psf1 = psf1 / psf1.sum()
+    return psf1, anchor
+
+
+def hsv_transform(img):
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+    # hsv[:, :, 0] = hsv[:, :, 0] * (0.2 + np.random.random() * 0.8)
+    hsv[:, :, 1] = hsv[:, :, 1] * (0.5 + np.random.random()*0.5)
+    hsv[:, :, 2] = hsv[:, :, 2] * (0.5 + np.random.random()*0.5)
+    img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+    return img
+
+
+def kth_diag_indices(array, k=0, n=-1):
+    """
+    第k个对角线索引
+    :param array: 输入二维矩阵
+    :param k: k=0 主对角线 k>0 右移 k<0 左移
+    :param n: n=-1, 对角线所有元素，否则只返回长为n的列表
+    :return: 索引
+    """
+    if n == -1:
+        rows, cols = np.diag_indices_from(array)
+    else:
+        rows, cols = np.diag_indices(n)
+    if k < 0:
+        return rows[-k:], cols[:k]
+    elif k > 0:
+        return rows[:-k], cols[k:]
+    else:
+        return rows, cols
+
+
+def slash_mask(mask, start_index, end_index, set_value=0.0, mod='l', length=-1):
+    """
+    制作斜条纹的蒙版
+    :param mask:
+    :param start_index:
+    :param end_index:
+    :param set_value:
+    :param mod: 左斜还是右斜
+    :param length: 长度
+    :return:
+    """
+    h, w = mask.shape[:2]
+    assert length <= min(h, w)
+    if mod == 'r':
+        mask = np.fliplr(mask)
+    for i in range(start_index, end_index+1):
+        mask[kth_diag_indices(mask, i, length)] = set_value
+    if mod == 'r':
+        mask = np.fliplr(mask)
+    return mask
+
+
+def line_mask(mask, start_index, end_index, set_value=0.0, mod='h', length=-1):
+    """
+    制作条纹蒙版
+    :param mask:
+    :param start_index:
+    :param end_index:
+    :param set_value:
+    :param mod: h 横 v 竖
+    :param length: 长度
+    :return:
+    """
+    h, w = mask.shape[:2]
+    if mod == 'h':
+        assert length <= w
+        assert 0 <= start_index < end_index < h
+        if length == -1 or length == w:
+            mask[start_index: end_index+1, :] = set_value
+        else:
+            left = random.randint(0, w-length-1)
+            right = left + length
+            mask[start_index: end_index + 1, left:right] = set_value
+    else:
+        assert length <= h
+        assert 0 <= start_index < end_index < w
+        if length == -1 or length == h:
+            mask[:, start_index: end_index + 1] = set_value
+        else:
+            top = random.randint(0, h - length - 1)
+            bottom = top + length
+            mask[top:bottom + 1, start_index: end_index] = set_value
+    return mask
+
+
+def read_binary_images(file_path):
+    img = np.fromfile(file_path, dtype=np.uint8)
+    img = np.reshape(img, (720, 1280, 3))
+    cv2.imshow('img', img)
+    cv2.waitKey(0)
+    print()
+
+
+
+if __name__ == '__main__':
+    read_binary_images('/Users/didi/Desktop/座次/pic/bgr1641010923242.bgr')