scbc_repos/ssat_sdk/picture/cie_luv.py

# -*- coding:utf-8 -*-
'''
颜色空间CIE LUV
'''
import os, sys, time
import cv2 as cv
import numpy as np
import colorsys
import math
class LUV():
    def __init__(self, stdimg = None, testimg=None):
        self.stdimg = stdimg
        self.testimg = testimg

    def setSTDImage(self, img):
        self.stdimg = img

    def setTestImage(self ,img):
        self.testimg = img

    '''
    计算单个指定区域的LUV平均值
    '''
    def getAverageLUV(self, img, area=None):
        if (area is None):
            area = (0, 0, img.shape[1], img.shape[0])  # mat shape:高、宽、颜色值个数
        # print "getAverageLUV", img.shape, area
        # print img.shape
        img = np.array(img)
        if len(area) == 4:
            x1 = area[0]
            y1 = area[1]
            x2 = area[2]
            y2 = area[3]
            img = img[y1:y2+1, x1:x2+1]
            # img = img[0:80, 0:100]
            # print img.shape
            total_L = 0
            total_U = 0
            total_V = 0
            pixel_total = img.shape[0] * img.shape[1]
            # print "pixel_total:",pixel_total

            for x in range(0, img.shape[0]):
                for y in range(0, img.shape[1]):
                    total_L = total_L + img[x][y][0]
                    total_U = total_U + img[x][y][1]
                    total_V = total_V + img[x][y][2]
            return total_L / pixel_total, total_U / pixel_total, total_V / pixel_total

        else:
            return None, None, None

    '''
    计算所有区域的LUV平均值
    '''
    def getMAverageLUV(self, img, areaList):
        total_L ,total_U, total_V, pixel_total = 0 ,0 ,0 ,0
        for area in areaList:
            L ,U ,V ,P = self.countTotalPLuv(img, area)
            total_L += L
            total_U += U
            total_V += V
            pixel_total += P
            # print "getMAverageLUV:",area,L,U,V,P
        if (pixel_total == 0):
            return None, None, None
        return total_L /pixel_total ,total_U /pixel_total ,total_V /pixel_total

    '''
    计算整张图片指定区域Luv颜色值总和。
    '''
    def countTotalPLuv(self, img, area=None):
        if (area is None):
            area = (0, 0, img.shape[1], img.shape[0])  # mat shape:高、宽、颜色值个数
        # print "getAverageLUV", img.shape, area
        # print img.shape
        img = np.array(img)
        if len(area) == 4:
            x1 = area[0]
            y1 = area[1]
            x2 = area[2]
            y2 = area[3]
            img = img[y1:y2 +1, x1:x2 +1]
            # img = img[0:80, 0:100]
            # print img.shape
            total_L = 0
            total_U = 0
            total_V = 0
            pixel_total = img.shape[0] * img.shape[1]
            # print "pixel_total:",pixel_total

            for x in range(0, img.shape[0]):
                for y in range(0, img.shape[1]):
                    total_L = total_L + img[x][y][0]
                    total_U = total_U + img[x][y][1]
                    total_V = total_V + img[x][y][2]
            return total_L ,total_U , total_V , pixel_total

        else:
            return None, None, None, None

    def getDiffLevel(self, stdLuv, testLuv):
        return self.getDiffUV(stdLuv, testLuv)

    def getDiffuvLevel(self, stdLuv, testLuv):
        diff = self.getDiffuv(stdLuv, testLuv)
        return diff[2]

    def getDiffluvLevel(self ,stdLuv, testLuv):
        return self.getDiffLUV(stdLuv, testLuv)

    def getDiffUV(self, stdLuv, testLuv):
        tmp = np.int32(testLuv[0]) - np.int32(stdLuv[0])
        lDiff = abs(tmp)
        tmp = np.int32(testLuv[1]) - np.int32(stdLuv[1])
        uDiff = abs(tmp)
        tmp = np.int32(testLuv[2]) - np.int32(stdLuv[2])
        vDiff = abs(tmp)
        lineDiff = np.sqrt(np.square(uDiff )+ np.square(vDiff))
        return uDiff ,vDiff ,lineDiff

    def getDiffuv(self, stdLuv, testLuv):
        return self.getDiffUV(stdLuv, testLuv)

    def getDiffLUV(self, stdLuv, testLuv):
        tmp = np.int32(testLuv[0]) - np.int32(stdLuv[0])
        lDiff = abs(tmp)
        tmp = np.int32(testLuv[1]) - np.int32(stdLuv[1])
        uDiff = abs(tmp)
        tmp = np.int32(testLuv[2]) - np.int32(stdLuv[2])
        vDiff = abs(tmp)
        lineDiff = np.sqrt(np.square(lDiff) + np.square(uDiff )+ np.square(vDiff))
        return lDiff, uDiff ,vDiff ,lineDiff

    '''
    # 描述：获取HSV平均值
    # 参数：
    # img_path：原图路径
    # area：要截图的区域
    # '''
    def getAverageHSV(self, img_path, area = None):
        img = cv.imread(img_path)
        if img is None:
            print u'Img is None, please check the img path！'
            return None, None, None

        # 求RGB;
        if (area is None):
            area = (0, 0, img.shape[1], img.shape[0])  # mat shape:高、宽、颜色值个数

        if len(area) == 4:
            x1 = area[0]
            y1 = area[1]
            x2 = area[2]
            y2 = area[3]
            img = img[y1:y2 +1, x1:x2 +1]
            # 宽，高;
            width, height = abs(x2 - x1), abs(y2 - y1)
            # print 'area=',area

            total_B = 0
            total_G = 0
            total_R = 0
            # pixel_total = hsv.shape[0] * hsv.shape[1]
            pixel_total = width * height
            if width == 0 or height == 0:
                print u'宽或高为0'
                return None, None, None
            # print 'width,hieght=',width,height,pixel_total

            for y in range(0, height):
                for x in range(0, width):
                    total_B += img[y][x][0]
                    total_G += img[y][x][1]
                    total_R += img[y][x][2]
            # print u't->b,g,r=',total_B,total_G,total_R
            b ,g ,r = total_B / pixel_total, total_G / pixel_total, total_R / pixel_total
            # print 'bgr=',b,g,r
            # 转成hsv;
            hsv = self.rgb2hsv(r ,g ,b)
            # print 'hsv-1',hsv
            # hsv = colorsys.rgb_to_hsv(r,g,b)
            # print 'hsv-2',hsv
            return [round(hsv[0 ] /2), round(hsv[1]), round(hsv[2])]
        else:
            return [None, None, None]

    # def rgb2hsv(self, r, g, b):
    #     r, g, b = r/255.0, g/255.0, b/255.0
    #     mx = max(r, g, b)
    #     mn = min(r, g, b)
    #     df = mx-mn
    #     if mx == mn:
    #         h = 0
    #     elif mx == r:
    #         h = (60 * ((g-b)/df) + 360) % 360
    #     elif mx == g:
    #         h = (60 * ((b-r)/df) + 120) % 360
    #     elif mx == b:
    #         h = (60 * ((r-g)/df) + 240) % 360
    #     if mx == 0:
    #         s = 0
    #     else:
    #         s = df/mx
    #     v = mx
    #     return h, s, v

    def rgb2hsv(self, r, g, b):
        r, g, b = r/ 255.0, g / 255.0, b / 255.0
        mx = max(r, g, b)
        mn = min(r, g, b)
        df = mx - mn
        # h值
        if mx == mn:
            h = 0
        elif mx == r and g >= b:
            h = (60 * (g - b) / df) % 360
        elif mx == r and g < b:
            h = (60 * ((g - b) / df) + 360) % 360
        elif mx == g:
            h = (60 * ((b - r) / df) + 120) % 360
        elif mx == b:
            h = (60 * ((r - g) / df) + 240) % 360
            # s值
        if mx == 0:
            s = 0
        else:
            s = df / mx * 255
        # v值
        v = mx * 255
        # 返回;
        return h, s, v


    def isUVSimilar(self, stdLuv, testLuv, diffLevel=20):
        udiff, vdiff, lineDiff = self.getDiffUV(stdLuv, testLuv)
        if lineDiff <= diffLevel:
            return True
        else:
            return False