151 lines
4.0 KiB
Python
151 lines
4.0 KiB
Python
import cv2;
|
||
import math;
|
||
import os
|
||
import numpy as np;
|
||
|
||
def DarkChannel(im,sz):
|
||
b,g,r = cv2.split(im)
|
||
dc = cv2.min(cv2.min(r,g),b);
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))
|
||
dark = cv2.erode(dc,kernel)
|
||
return dark
|
||
|
||
def AtmLight(im,dark):
|
||
[h,w] = im.shape[:2]
|
||
imsz = h*w
|
||
numpx = int(max(math.floor(imsz/1000),1))
|
||
darkvec = dark.reshape(imsz);
|
||
imvec = im.reshape(imsz,3);
|
||
|
||
indices = darkvec.argsort();
|
||
indices = indices[imsz-numpx::]
|
||
|
||
atmsum = np.zeros([1,3])
|
||
for ind in range(1,numpx):
|
||
atmsum = atmsum + imvec[indices[ind]]
|
||
|
||
A = atmsum / numpx;
|
||
return A
|
||
|
||
def TransmissionEstimate(im,A,sz):
|
||
omega = 0.95;
|
||
im3 = np.empty(im.shape,im.dtype);
|
||
|
||
for ind in range(0,3):
|
||
im3[:,:,ind] = im[:,:,ind]/A[0,ind]
|
||
|
||
transmission = 1 - omega*DarkChannel(im3,sz);
|
||
return transmission
|
||
|
||
def Guidedfilter(im,p,r,eps):
|
||
mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r));
|
||
mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r));
|
||
mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r));
|
||
cov_Ip = mean_Ip - mean_I*mean_p;
|
||
|
||
mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r));
|
||
var_I = mean_II - mean_I*mean_I;
|
||
|
||
a = cov_Ip/(var_I + eps);
|
||
b = mean_p - a*mean_I;
|
||
|
||
mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r));
|
||
mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r));
|
||
|
||
q = mean_a*im + mean_b;
|
||
return q;
|
||
|
||
def TransmissionRefine(im,et):
|
||
gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY);
|
||
gray = np.float64(gray)/255;
|
||
r = 60;
|
||
eps = 0.0001;
|
||
t = Guidedfilter(gray,et,r,eps);
|
||
|
||
return t;
|
||
|
||
def Recover(im,t,A,tx = 0.1):
|
||
res = np.empty(im.shape,im.dtype);
|
||
t = cv2.max(t,tx);
|
||
|
||
for ind in range(0,3):
|
||
res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]
|
||
|
||
return res
|
||
|
||
def getFileList(dir,Filelist, ext=None):
|
||
"""
|
||
获取文件夹及其子文件夹中文件列表
|
||
输入 dir：文件夹根目录
|
||
输入 ext: 扩展名
|
||
返回： 文件路径列表
|
||
"""
|
||
newDir = dir
|
||
if os.path.isfile(dir):
|
||
if ext is None:
|
||
Filelist.append(dir)
|
||
else:
|
||
if ext in dir[-3:]:
|
||
Filelist.append(dir)
|
||
|
||
elif os.path.isdir(dir):
|
||
for s in os.listdir(dir):
|
||
newDir=os.path.join(dir,s)
|
||
getFileList(newDir, Filelist, ext)
|
||
|
||
return Filelist
|
||
|
||
if __name__ == '__main__':
|
||
import sys
|
||
|
||
sz = 10 # 窗口函数
|
||
tx = 0.2 # 传输图最小值
|
||
|
||
try:
|
||
local_dir = sys.argv[1]
|
||
except:
|
||
local_dir = './image/'
|
||
|
||
try:
|
||
sz = int(sys.argv[2])
|
||
except:
|
||
sz = sz
|
||
|
||
try:
|
||
tx = float(sys.argv[3])
|
||
except:
|
||
tx = tx
|
||
|
||
def nothing(*argv):
|
||
pass
|
||
# 检索文件
|
||
src_img_folder = os.path.join(local_dir, 'src')
|
||
|
||
imglist = getFileList(src_img_folder, [], '')
|
||
print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
|
||
|
||
for imgpath in imglist:
|
||
imgname= os.path.splitext(os.path.basename(imgpath))[0]
|
||
# cv2.IMREAD_GRAYSCALE / cv2.IMREAD_COLOR ： 加载灰色 / 彩色图像
|
||
src = cv2.imread(imgpath, cv2.IMREAD_COLOR)
|
||
# 通道归一滑
|
||
I = src.astype('float64')/255;
|
||
# 暗通道图像
|
||
dark = DarkChannel(I,sz=sz);
|
||
# ？？？
|
||
A = AtmLight(I,dark);
|
||
te = TransmissionEstimate(I,A,sz = sz);
|
||
t = TransmissionRefine(src,te);
|
||
print(np.shape(src))
|
||
print(np.shape(dark))
|
||
print(np.shape(t))
|
||
J = Recover(I,t,A,tx=tx); # tx传输图的最小值，用于避免过度曝光
|
||
|
||
dark_imgdir = os.path.join(local_dir, 'dark/')
|
||
trans_imgdir = os.path.join(local_dir, 'trans/')
|
||
result_imgdir = os.path.join(local_dir, 'result/')
|
||
cv2.imwrite(dark_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_dark.png",dark*255);
|
||
cv2.imwrite(trans_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_t.png", t*255);
|
||
cv2.imwrite(result_imgdir + imgname + "_" + str(sz) + "_" + str(tx) + "_result.png",J*255);
|
||
|