scbc_repos/ssat_sdk/picture/pic_segment.py

# -*- coding:utf-8 -*-
import os,sys,time
import cv2 as cv
import numpy as np
from ssat_sdk.utils import LoggingUtil
from ssat_sdk.picture import image_util

class PicSegment():
    def __init__(self):
        pass


    '''
    在传入的图片上，查找指定bgr颜色框矩形区域，返回矩形区域的坐标
    :param picPath:图片路径
    :param bgr：颜色值
    :return result, contourRect：result=True/False, contourRect:矩形区域
    '''
    def findAreaByColor(self, srcImg, cropArea, bgr, offset=20, minPeri=0, maxPeri=0, minArea=0, maxArea=0,
                                morphology=[]):
        # 获取rgb二值图;
        img_binary, cropArea = self.getBGRBinaryImage(srcImg, cropArea, bgr, offset)
        if img_binary is None:
            return False, None

        # 查找符合要求的外轮廓;
        # opencv3开始，有3个返回值：img, countours, hierarchy
        contours = cv.findContours(img_binary, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)[1]
        # 是否查找最大轮廓;
        bFindMaxContour = False
        # 如果周长、面积都一样，只返回最大轮廓;
        if minPeri == maxPeri and minArea == maxArea:
            bFindMaxContour = True

        contourRect = []
        maxContour = None  # 最大轮廓;
        tempArea = 0
        # 过滤掉不符合要求的轮廓;
        x, y, w, h, result = 0, 0, 0, 0, False
        for cnt in contours:
            # 面积;
            area = cv.contourArea(cnt)
            # 周长;
            perimeter = cv.arcLength(cnt, True)
            # print "area=%ld, perimeter=%ld"%(area, perimeter)
            if bFindMaxContour == True:
                # 获取最大轮廓;
                if tempArea < area:
                    tempArea = area
                    maxContour = cnt
            else:
                # print area,perimeter
                # 获取满足条件的值;
                if area >= minArea and area <= maxArea and perimeter >= minPeri and perimeter <= maxPeri:
                    # 直边矩形,(x,y）为矩形左上角的坐标，（w,h）是矩形的宽和高;
                    x, y, w, h = cv.boundingRect(cnt)
                    # 在区域内截图判断;
                    if cropArea is not None:
                        x = x + cropArea[0]
                        y = y + cropArea[1]
                    contourRect = [x, y, x + w, y + h]
                    result = True
                    break
                # endif
        # endfor

        if bFindMaxContour == True and maxContour is not None:
            result = True
            # 最大轮廓的矩形坐标;
            x, y, w, h = cv.boundingRect(maxContour)
            # 在区域内截图判断;
            if cropArea is not None:
                x = x + cropArea[0]
                y = y + cropArea[1]
            contourRect = [x, y, x + w, y + h]

        # 只返回结果、轮廓坐标;
        return result, contourRect

    '''
    不做任何图片补足处理的颜色区域处理，和二值化图片获取
    '''
    def getBGRBinaryImage(self, srcImg, cropArea, bgr, offset=20):
        # 读取图片;
        if srcImg is None:
            return None,None
        # 如果截图区域空，使用原图大小;
        if cropArea is not None or cropArea.__len__() == 0:
            cropArea = [0, 0, srcImg.shape[1], srcImg.shape[0]]

        img_crop = None
        # 只保留区域部分;
        try:
            img_crop = image_util.cutImage(srcImg, cropArea)
        except:
            img_crop = srcImg
        subImg = self.subImgBgr(img_crop, bgr)
        cv.imwrite(r'd:\\subImg.png', subImg)
        subImg_wh1 = np.where((subImg > -offset) & (subImg < offset), np.uint8(0), np.uint8(255))
        cv.imwrite(r'd:\\subImg_wh1.png', subImg_wh1)
        # 将图处转成灰阶
        img_gray = cv.cvtColor(subImg_wh1, cv.COLOR_BGR2GRAY)
        # 反转灰度值，以便适合查找轮廓;
        img_gray = np.where((img_gray < -offset) | (img_gray > offset), np.uint8(0), np.uint8(255))
        # 输出结果图片;
        cv.imwrite(r'd:\\img_gray.png', img_gray)
        # 使用灰度图求得二值化图;
        thresh = cv.threshold(img_gray, 1, 255, cv.THRESH_BINARY)[1]
        # 返回二值化图;
        return thresh, cropArea

    '''
    在原有图片，减去固定的BGR颜色值，用于后续的颜色筛选.返回减去bgr的图片。
    '''
    def subImgBgr(self, bgrImg, bgr):
        imgArr = np.asarray(bgrImg).astype(np.int16)
        bgrArr = np.asarray(bgr).astype(np.int16)
        # 矩阵相减;
        subImg = imgArr - bgrArr
        return subImg