整合去雾网页工具
This commit is contained in:
admin
208
DehazeNet/DehazeNet.py
Normal file
208
DehazeNet/DehazeNet.py
Normal file
@@ -0,0 +1,208 @@
|
||||
import sys,os
|
||||
import caffe
|
||||
import numpy as np
|
||||
import cv2
|
||||
import math
|
||||
|
||||
def EditFcnProto(templateFile, height, width):
|
||||
with open(templateFile, 'r') as ft:
|
||||
template = ft.read()
|
||||
outFile = 'DehazeNetFcn.prototxt'
|
||||
with open(outFile, 'w') as fd:
|
||||
fd.write(template.format(height_15=height+15, width_15=width+15,
|
||||
height_11=height+11, width_11=width+11))
|
||||
|
||||
def TransmissionEstimate(im_path, height, width):
|
||||
caffe.set_mode_cpu()
|
||||
|
||||
# Define a safe tile size to prevent INT_MAX overflow (approx 512x512 is safe)
|
||||
SAFE_TILE_SIZE = 512
|
||||
|
||||
# Use tiling if the image is larger than the safe size
|
||||
if height > SAFE_TILE_SIZE or width > SAFE_TILE_SIZE:
|
||||
print(f"Image size ({width}x{height}) is large. Using tiling to avoid memory overflow...")
|
||||
|
||||
# Determine effective tile size (cannot be larger than image)
|
||||
tile_h = min(height, SAFE_TILE_SIZE)
|
||||
tile_w = min(width, SAFE_TILE_SIZE)
|
||||
|
||||
# Generate prototxt for the TILE size, not the full image size
|
||||
EditFcnProto('DehazeFcnTemplate.prototxt', tile_h, tile_w)
|
||||
|
||||
# Load networks
|
||||
net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
|
||||
net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
|
||||
net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
|
||||
net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
|
||||
|
||||
# Load and pad image
|
||||
im = caffe.io.load_image(im_path)
|
||||
npad = ((7,8), (7,8), (0,0))
|
||||
im_padded = np.pad(im, npad, 'symmetric')
|
||||
|
||||
transmission = np.zeros((height, width))
|
||||
|
||||
# Setup transformer for the tile size
|
||||
transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
|
||||
transformers.set_transpose('data', (2,0,1))
|
||||
transformers.set_channel_swap('data', (2,1,0))
|
||||
|
||||
# Process in tiles
|
||||
for h in range(0, height, tile_h):
|
||||
for w in range(0, width, tile_w):
|
||||
# Calculate start/end to handle edges/overlap
|
||||
# If we are at the end, shift back to ensure we feed a full tile
|
||||
h_start = min(h, height - tile_h)
|
||||
w_start = min(w, width - tile_w)
|
||||
|
||||
# Extract patch from PADDED image
|
||||
# Network expects input size of (Tile + 15), so we slice accordingly
|
||||
patch = im_padded[h_start : h_start + tile_h + 15, w_start : w_start + tile_w + 15, :]
|
||||
|
||||
# Forward pass
|
||||
out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', patch-0.2)]))
|
||||
|
||||
# Reshape output
|
||||
block_trans = np.reshape(out['ip1-conv'], (tile_h, tile_w))
|
||||
|
||||
# Assign to result buffer
|
||||
transmission[h_start : h_start + tile_h, w_start : w_start + tile_w] = block_trans
|
||||
|
||||
return transmission
|
||||
|
||||
else:
|
||||
# Original logic for small images
|
||||
EditFcnProto('DehazeFcnTemplate.prototxt', height, width)
|
||||
net = caffe.Net('DehazeNet.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
|
||||
net_full_conv = caffe.Net('DehazeNetFcn.prototxt', 'DehazeNet.caffemodel', caffe.TEST)
|
||||
net_full_conv.params['ip1-conv'][0].data.flat = net.params['ip1'][0].data.flat
|
||||
net_full_conv.params['ip1-conv'][1].data[...] = net.params['ip1'][1].data
|
||||
im = caffe.io.load_image(im_path)
|
||||
npad = ((7,8), (7,8), (0,0))
|
||||
im = np.pad(im, npad, 'symmetric')
|
||||
transformers = caffe.io.Transformer({'data': net_full_conv.blobs['data'].data.shape})
|
||||
transformers.set_transpose('data', (2,0,1))
|
||||
transformers.set_channel_swap('data', (2,1,0))
|
||||
out = net_full_conv.forward_all(data=np.array([transformers.preprocess('data', im-0.2)]))
|
||||
transmission = np.reshape(out['ip1-conv'], (height,width))
|
||||
return transmission
|
||||
|
||||
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,1)
|
||||
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 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__':
|
||||
if not len(sys.argv) == 2:
|
||||
print ('Usage: python DeHazeNet.py haze_img_path')
|
||||
exit()
|
||||
else:
|
||||
im_path = sys.argv[1]
|
||||
|
||||
# 检索文件
|
||||
src_img_folder = os.path.join(im_path, 'src')
|
||||
imglist = getFileList(src_img_folder, [], '')
|
||||
print('本次执行检索到 '+str(len(imglist))+' 张图像\n')
|
||||
for img_path in imglist:
|
||||
imgname= os.path.splitext(os.path.basename(img_path))[0]
|
||||
src = cv2.imread(img_path)
|
||||
|
||||
height = src.shape[0]
|
||||
width = src.shape[1]
|
||||
|
||||
# Note: EditFcnProto is also called inside TransmissionEstimate if tiling is used
|
||||
# We call it here for initialization but it may be overwritten.
|
||||
templateFile = 'DehazeFcnTemplate.prototxt'
|
||||
EditFcnProto(templateFile, height, width)
|
||||
print("-"*5, ' 完成EditFcnProto ',"-"*5)
|
||||
|
||||
I = src/255.0
|
||||
dark = DarkChannel(I,15)
|
||||
A = AtmLight(I,dark)
|
||||
|
||||
te = TransmissionEstimate(img_path, height, width)
|
||||
t = TransmissionRefine(src,te)
|
||||
J = Recover(I,t,A,0.1)
|
||||
|
||||
print("Finsh All the operation")
|
||||
Trans_Esti_imgdir = os.path.join(im_path, 'Trans_Esti/')
|
||||
if not os.path.exists(Trans_Esti_imgdir): os.makedirs(Trans_Esti_imgdir)
|
||||
print(Trans_Esti_imgdir + imgname + "_Trans_Esti.png")
|
||||
cv2.imwrite(Trans_Esti_imgdir + imgname + "_Trans_Esti.png",te*255);
|
||||
|
||||
Trans_Refine_imgdir = os.path.join(im_path, 'Trans_Refine/')
|
||||
if not os.path.exists(Trans_Refine_imgdir): os.makedirs(Trans_Refine_imgdir)
|
||||
print(Trans_Refine_imgdir + imgname + "_Trans_Refine.png")
|
||||
cv2.imwrite(Trans_Refine_imgdir + imgname + "_Trans_Refine.png",t*255);
|
||||
|
||||
result_imgdir = os.path.join(im_path, 'result/')
|
||||
if not os.path.exists(result_imgdir): os.makedirs(result_imgdir)
|
||||
print(result_imgdir + imgname + "_result.png")
|
||||
cv2.imwrite(result_imgdir + imgname + "_result.png",J*255);
|
||||
Reference in New Issue
Block a user