暗通道先验对单图像去雾 Single Image Haze Removal Using Dark Channel Prior

原文发表在PAMI2011,一作何凯明。

# -*- coding: utf-8 -*-
import cv2
import numpy as np
import argparse


def dark_channel(img, size = 15):
    r, g, b = cv2.split(img)
    min_img = cv2.min(r, cv2.min(g, b))#求出三通道最低值
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))#这个函数的第一个参数表示内核的形状,有三种形状可以选择。矩形:MORPH_RECT;交叉形:MORPH_CROSS;椭圆形:MORPH_ELLIPSE;
    #获得15*15的锚框
    dc_img = cv2.erode(min_img,kernel)
    return dc_img


def get_atmo(img, percent = 0.001):#获取大气A值,具体做法是在暗通道中获取亮度为前0.1%的像素,再在这些像素中寻找亮度最高的值作为A值。
    mean_perpix = np.mean(img, axis = 2).reshape(-1)
    mean_topper = mean_perpix[:int(img.shape[0] * img.shape[1] * percent)]
    return np.mean(mean_topper)


def get_trans(img, atom, w = 0.95):
    x = img / atom
    t = 1 - w * dark_channel(x, 15)
    return t


def guided_filter(p, i, r, e):
    """
    :param p: input image
    :param i: guidance image
    :param r: radius
    :param e: regularization
    :return: filtering output q
    """
    #1
    mean_I = cv2.boxFilter(i, cv2.CV_64F, (r, r))
    mean_p = cv2.boxFilter(p, cv2.CV_64F, (r, r))
    corr_I = cv2.boxFilter(i * i, cv2.CV_64F, (r, r))
    corr_Ip = cv2.boxFilter(i * p, cv2.CV_64F, (r, r))
    #2
    var_I = corr_I - mean_I * mean_I
    cov_Ip = corr_Ip - mean_I * mean_p
    #3
    a = cov_Ip / (var_I + e)
    b = mean_p - a * mean_I
    #4
    mean_a = cv2.boxFilter(a, cv2.CV_64F, (r, r))
    mean_b = cv2.boxFilter(b, cv2.CV_64F, (r, r))
    #5
    q = mean_a * i + mean_b
    return q


def dehaze(path, output = None):
    im = cv2.imread(path)
    img = im.astype('float64') / 255
    img_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY).astype('float64') / 255

    atom = get_atmo(img)
    trans = get_trans(img, atom)
    trans_guided = guided_filter(trans, img_gray, 20, 0.0001)
    trans_guided = cv2.max(trans_guided, 0.25)

    result = np.empty_like(img)
    for i in range(3):
        result[:, :, i] = (img[:, :, i] - atom) / trans_guided + atom

    cv2.imshow("source",img)
    cv2.imshow("result", result)
    cv2.waitKey()
    if output is not None:
        cv2.imwrite(output, result * 255)


parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input')
parser.add_argument('-o', '--output')
args = parser.parse_args()


if __name__ == '__main__':
    if args.input is None:
        dehaze('3.jpg')
    else:
        dehaze(args.input, args.output)

推荐参考,比我写的清楚得多,而且实验非常扎实:
https://www.cnblogs.com/Imageshop/p/3281703.html

你可能感兴趣的:(计算机视觉,opencv)