Histogram equalization and specialization

 

from PIL import Image
import os


class Img:
    def __init__(self, image):
        self.img = image

    def getPr(self):
        # 获取各灰度的概率
        img = self.img
        dic = {}

        colors = img.getcolors()

        totalPixels = img.width*img.height

        for counts, r in colors:
            dic[r] = counts/totalPixels

        return dic

    def getEqualization(self):
        # 计算均衡字典{inputPix : outputPix}
        img = self.img
        sumPr = 0

        dic_equal = self.getPr()

        for r in dic_equal:
            sumPr += dic_equal[r]
            dic_equal[r] = round(sumPr * 255)

        return dic_equal

    def getEqualReverse(self):
        # 计算反均衡字典{outputPix : inputPix}
        dic_equal = self.getEqualization()
        dic_reverse = {}

        for r in dic_equal:
            if not dic_equal[r] in dic_reverse:
                dic_reverse[dic_equal[r]] = r

        return dic_reverse

    def getEqualizedImg(self):
        # 计算均衡后的图像
        img = self.img
        dic_equal = self.getEqualization()

        equalizedImg = Image.new(
            img.mode, (img.width, img.height))

        for x in range(img.width):
            for y in range(img.height):
                r = img.getpixel((x, y))
                equalizedImg.putpixel((x, y), dic_equal[r])

        return Img(equalizedImg)

    def getSpecification(self, targetImg):
        # 计算规定化后的图像
        img = self.img

        dic = self.getEqualization()
        dic2 = targetImg.getEqualReverse()

        for r in dic:
            if dic[r] in dic2:
                dic[r] = dic2[dic[r]]
            else:
                s = dic[r]
                i, j = 1, 1
                while (not s + i in dic2):
                    if s + i > 255:
                        i = 256
                        break
                    i += 1
                while (not s - j in dic2):
                    if s - j < 0:
                        j = 256
                        break
                    j += 1
                if i < j:
                    dic[r] = dic2[s + i]
                else:
                    dic[r] = dic2[s - j]

        specializedImg = Image.new(
            img.mode, (img.width, img.height))

        for x in range(img.width):
            for y in range(img.height):
                r = img.getpixel((x, y))
                specializedImg.putpixel((x, y), dic[r])

        return Img(specializedImg)


whiteImgName = '322_sli128_T1.png'  # 目标图像
sourceFolderName = 'NoEyeData'
savingFolderName = 'NoEyeData_s'

if not os.path.exists('./' + savingFolderName):
    os.mkdir('./' + savingFolderName)

filePath = './' + sourceFolderName
fileList = os.listdir(filePath)

white = Image.open(filePath + '/'+whiteImgName)
white.save('./'+savingFolderName+'/'+whiteImgName)
white = Img(white)

for i in fileList:
    if i == whiteImgName:
        continue
    black = Image.open(filePath+'/'+i)
    black = Img(black)

    b2w = black.getSpecification(white)
    b2w.img.save('./'+savingFolderName+'/'+i)