EMC/app/api/model/model.py

import cv2
import numpy as np
import argparse
import onnxruntime as ort
from pathlib import Path
from tqdm import tqdm
import time
import sys
import os
import requests
from utils.utils import send_notifycation

def getFile(ruleFile):
    if getattr(sys, 'frozen', False):
        absPath = os.path.dirname(os.path.abspath(sys.executable))
    elif __file__:
        absPath = os.path.dirname(os.path.abspath(__file__))
    else:
        absPath = ''
    return os.path.join(absPath,ruleFile)

def checkModel(model):
    """
    【功能】检查模型是否存在
    【参数】
        model: str   必选,模型名称
    【输入/输出】
        True: 模型存在
        False: 模型不存在
    """
    # 先检查本地是否存在该文件模型，不存在则查看网络数据库是否存在该模型，存在则下载然后返回该模型
    if not os.path.exists(getFile(model)):
            send_notifycation('模型不存在，正在下载模型...')
            url = 'https://gitee.com/JiangNoah/emc/raw/master/model.json'
            r = requests.get(url)
            model_dict = r.json()# model_dict = {models:[{'name':model,'url':url}]}
            #检查是否存在该模型的name
            models = model_dict["models"]
            url_ = ''
            for m in models:
                if m["name"] == model:
                    url_ = m["url"]
                    break
            if url_ == '':
                send_notifycation('模型不存在')
                return getFile(model)
            else:
                send_notifycation('模型下载中...')
                r = requests.get(url_)
                with open(getFile(model), 'wb') as f:
                    f.write(r.content)
                send_notifycation('模型下载完成')
                return getFile(model)
    else:
        return getFile(model)


class PicoDet():
    def __init__(self,
                 model_pb_path,
                 label_path,
                 prob_threshold=0.4,
                 iou_threshold=0.3):
        self.classes = list(
            map(lambda x: x.strip(), open(getFile(label_path), 'r').readlines()))
        self.num_classes = len(self.classes)
        self.prob_threshold = prob_threshold
        self.iou_threshold = iou_threshold
        self.mean = np.array(
            [103.53, 116.28, 123.675], dtype=np.float32).reshape(1, 1, 3)
        self.std = np.array(
            [57.375, 57.12, 58.395], dtype=np.float32).reshape(1, 1, 3)
        so = ort.SessionOptions()
        so.log_severity_level = 3
        self.net = ort.InferenceSession(checkModel(model_pb_path), so)
        inputs_name = [a.name for a in self.net.get_inputs()]
        inputs_shape = {
            k: v.shape
            for k, v in zip(inputs_name, self.net.get_inputs())
        }
        self.input_shape = inputs_shape['image'][2:]

    def _normalize(self, img):
        img = img.astype(np.float32)
        img = (img / 255.0 - self.mean / 255.0) / (self.std / 255.0)
        return img

    def resize_image(self, srcimg, keep_ratio=False):
        top, left, newh, neww = 0, 0, self.input_shape[0], self.input_shape[1]
        origin_shape = srcimg.shape[:2]
        im_scale_y = newh / float(origin_shape[0])
        im_scale_x = neww / float(origin_shape[1])
        img_shape = np.array([
            [float(self.input_shape[0]), float(self.input_shape[1])]
        ]).astype('float32')
        scale_factor = np.array([[im_scale_y, im_scale_x]]).astype('float32')

        if keep_ratio and srcimg.shape[0] != srcimg.shape[1]:
            hw_scale = srcimg.shape[0] / srcimg.shape[1]
            if hw_scale > 1:
                newh, neww = self.input_shape[0], int(self.input_shape[1] /
                                                      hw_scale)
                img = cv2.resize(
                    srcimg, (neww, newh), interpolation=cv2.INTER_AREA)
                left = int((self.input_shape[1] - neww) * 0.5)
                img = cv2.copyMakeBorder(
                    img,
                    0,
                    0,
                    left,
                    self.input_shape[1] - neww - left,
                    cv2.BORDER_CONSTANT,
                    value=0)  # add border
            else:
                newh, neww = int(self.input_shape[0] *
                                 hw_scale), self.input_shape[1]
                img = cv2.resize(
                    srcimg, (neww, newh), interpolation=cv2.INTER_AREA)
                top = int((self.input_shape[0] - newh) * 0.5)
                img = cv2.copyMakeBorder(
                    img,
                    top,
                    self.input_shape[0] - newh - top,
                    0,
                    0,
                    cv2.BORDER_CONSTANT,
                    value=0)
        else:
            img = cv2.resize(
                srcimg, self.input_shape, interpolation=cv2.INTER_AREA)

        return img, img_shape, scale_factor

    def get_color_map_list(self, num_classes):
        color_map = num_classes * [0, 0, 0]
        for i in range(0, num_classes):
            j = 0
            lab = i
            while lab:
                color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
                color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
                color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
                j += 1
                lab >>= 3
        color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
        return color_map

    def detect(self, srcimg, show_result=False):
        img, im_shape, scale_factor = self.resize_image(srcimg)
        img = self._normalize(img)

        blob = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)

        inputs_dict = {
            'im_shape': im_shape,
            'image': blob,
            'scale_factor': scale_factor
        }
        inputs_name = [a.name for a in self.net.get_inputs()]
        net_inputs = {k: inputs_dict[k] for k in inputs_name}

        outs = self.net.run(None, net_inputs)

        outs = np.array(outs[0])
        expect_boxes = (outs[:, 1] > 0.5) & (outs[:, 0] > -1)
        np_boxes = outs[expect_boxes, :]

        color_list = self.get_color_map_list(self.num_classes)
        clsid2color = {}

        result = []
        for i in range(np_boxes.shape[0]):
            classid, conf = int(np_boxes[i, 0]), np_boxes[i, 1]
            result.append({
                'classid': self.classes[classid],
                'conf': str(round(conf, 3)),
            })

        if(show_result):
            for i in range(np_boxes.shape[0]):
                classid, conf = int(np_boxes[i, 0]), np_boxes[i, 1]
                xmin, ymin, xmax, ymax = int(np_boxes[i, 2]), int(np_boxes[
                    i, 3]), int(np_boxes[i, 4]), int(np_boxes[i, 5])

                if classid not in clsid2color:
                    clsid2color[classid] = color_list[classid]
                color = tuple(clsid2color[classid])

                cv2.rectangle(
                    srcimg, (xmin, ymin), (xmax, ymax), color, thickness=2)
                print(self.classes[classid] + ': ' + str(round(conf, 3)))
                cv2.putText(
                    srcimg,
                    self.classes[classid] + ':' + str(round(conf, 3)), (xmin,
                                                                        ymin - 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.8, (0, 255, 0),
                    thickness=2)

            return [result,srcimg]
        else:
            
            return result

    def detect_folder(self, img_fold, result_path):
        img_fold = Path(img_fold)
        result_path = Path(result_path)
        result_path.mkdir(parents=True, exist_ok=True)

        img_name_list = filter(
            lambda x: str(x).endswith(".png") or str(x).endswith(".jpg"),
            img_fold.iterdir(), )
        img_name_list = list(img_name_list)
        print(f"find {len(img_name_list)} images")

        for img_path in tqdm(img_name_list):
            img = cv2.imread(str(img_path))

            srcimg = self.detect(img,True)
            save_path = str(result_path / img_path.name.replace(".png", ".jpg"))
            cv2.imwrite(save_path, srcimg)
    
    def detect_img(self, img,show_result=False):
        # img = cv2.imread(img)
        # img_path = Path(img_path)
        # print(f"find {img_path} images")
        
        
        if(show_result):
            [result,srcimg] = self.detect(img,show_result)
            return [result,srcimg]
        else:
            result = self.detect(img,show_result)
            return result
            # result_path = Path(result_path)
            # result_path.mkdir(parents=True, exist_ok=True)
            # save_path = str(result_path / "result.jpg")
            # cv2.imwrite(save_path, srcimg)
        #     return srcimg
        # else:

    def crop_2_224(img):
        height=len(img)
        width=len(img[0])
        if(height>224 and width>224):
            y0 = height/2
            x0 = width/2
            x1 = x0-112
            y1 = y0-112
            x2 = x0+112
            y2 = y0+112
            img = img[y1:y2, x1:x2]
        return img