CheckCardBoxAndDefectServer/app/utils/create_predict_result.py

import numpy as np
import cv2
from pathlib import Path
import time

from app.utils.data_augmentation import LetterBox


def point_mapTo_originImg(originImgSize, imgSize_train, now_point):
    letterBox = LetterBox(imgSize_train)
    rect_dict = letterBox.get_offset(originImgSize)

    x_ratio = originImgSize['width'] * 1.0 / rect_dict['width']
    y_ratio = originImgSize['height'] * 1.0 / rect_dict['height']

    new_y = round((now_point[1] - rect_dict['y']) * y_ratio)
    new_x = round((now_point[0] - rect_dict['x']) * x_ratio)

    new_point = [new_x, new_y]

    return new_point


def create_result_singleImg(segClassDict, now_ansImgDict, originImgSize, imgSize_train, confidence=0.5):
    """ansImg_list.append({"ans_img":ansImg,"probs":probs,"file_name":file_name})"""

    label = now_ansImgDict['ans_img']
    probs = now_ansImgDict['probs']
    file_name = now_ansImgDict['file_name']

    # confidence = 0.5
    assert confidence > 0, "置信度必须大于0"
    # label[label < confidence] = 0

    per_result = {}
    per_result['num'] = 0
    per_result['cls'] = []
    per_result['names'] = []
    per_result['conf'] = []
    per_result['shapes'] = []

    # 背景肯定不需要提取
    excludeClassList = ['___background___']

    # 遍历每个分类
    for key, val in segClassDict.items():

        if val not in excludeClassList:

            now_class_num = int(key)
            now_prob_img = probs[now_class_num]

            # 10、将对应分类的图片弄白，其余全黑
            imgZero = np.zeros(label.shape, dtype=np.uint8)
            imgZero[label == (now_class_num)] = 255

            # 20、检测轮廓
            contours, _ = cv2.findContours(imgZero, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            # drawnContourImg = np.zeros((imgZero.shape[0], imgZero.shape[1], 3), dtype=np.uint8)
            # 把轮廓画出来了，可能label就是对应的轮廓
            # drawnContourImg = cv2.drawContours(drawnContourImg, contours, -1, (255, 0, 0), 2)

            # 30、获取shape数据
            if len(contours):
                maxCntArea = 5

                # 找最大轮廓面积对应的index
                for index in range(len(contours)):
                    nowPntList = []
                    # 面积
                    contourArea = cv2.contourArea(contours[index])
                    # 面积要大于5
                    if contourArea > maxCntArea:

                        # 计算置信度
                        # 创建一个全为零的掩码图像，大小和原图像一样
                        mask = np.zeros_like(imgZero)

                        # 绘制第一个轮廓（假设我们只关心第一个轮廓）
                        cv2.drawContours(mask, contours, index, (255), thickness=cv2.FILLED)

                        # plt.imshow(now_prob_img, cmap='gray')
                        # plt.show()

                        # plt.imshow(mask, cmap='gray')
                        # plt.show()

                        # 计算轮廓区域的均值
                        mean_val = cv2.mean(now_prob_img, mask=mask)
                        now_conf = mean_val[0]

                        externalContour = contours[index]
                        pointNum = len(externalContour)
                        # print(pointNum)

                        for i in range(pointNum):
                            nowPoint = externalContour[i]
                            nowPoint_list = nowPoint[0].tolist()
                            # 点要放到到原图上面
                            new_point = point_mapTo_originImg(originImgSize, imgSize_train, nowPoint_list)

                            nowPntList.append(new_point)

                        pre_defect = {}
                        pre_defect['class_num'] = int(key)
                        pre_defect['label'] = str(val)
                        pre_defect['probability'] = now_conf
                        pre_defect['points'] = nowPntList

                        if now_conf >= confidence:
                            per_result['cls'].append(int(key))
                            per_result['names'].append(str(val))
                            per_result['conf'] = now_conf
                            per_result['shapes'].append(pre_defect)

    per_result['num'] = len(per_result['shapes'])
    return per_result