在对一个文件夹里面的图像归一化的时候难免会遇到要判断文件夹的内容是文件还是目录的情况,由于使用递归函数太繁琐,于是使用os.walk来做比较方便。
在代码里的image_normallize这个函数里面有一个area用来控制寻找通过面积过滤找到的轮廓。
import os
import cv2
import numpy as np
def rotate_bound(image, angle):
(h, w) = image.shape[:2]
(cX, cY) = (w / 2, h / 2)
M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
cos = np.abs(M[0, 0])
sin = np.abs(M[0, 1])
nW = int((h * sin) + (w * cos))
nH = int((h * cos) + (w * sin))
M[0, 2] += (nW / 2) - cX
M[1, 2] += (nH / 2) - cY
return cv2.warpAffine(image, M, (nW, nH))
def image_normalize(img_path, save_folder, filename, gray_thresh, cut_wh, min_area, max_area):
ori_image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)
ih, iw = ori_image.shape[:2]
gray_image = cv2.cvtColor(ori_image, cv2.COLOR_BGR2GRAY)
_, bi_image = cv2.threshold(gray_image, gray_thresh, 255, cv2.THRESH_BINARY)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (21, 21))
bi_image = cv2.dilate(bi_image, kernel)
counts, _ = cv2.findContours(bi_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
x, y, w, h, t = 0, 0, 0, 0, 90
for j in counts:
area = cv2.contourArea(j)
print("area:", area)
if max_area > area > min_area: # 面积过滤
rect = cv2.minAreaRect(j)
(x, y), (w, h), t = rect
if t > 45:
ori_image = rotate_bound(ori_image, 90 - t)
else:
ori_image = rotate_bound(ori_image, -t)
h1, w1 = ori_image.shape[:2]
offset_h, offset_w = h1 - ih, w1 - iw
new_image = ori_image[int(y - cut_wh[1]//2 + offset_h / 2):int(y + cut_wh[1]//2 + offset_h / 2),
int(x - cut_wh[0]//2 + offset_w / 2):int(x + cut_wh[0]//2 + offset_w / 2)]
try:
cv2.imencode(".jpg", new_image)[1].tofile(r"%s\%s" % (save_folder, filename))
print("save %s\%s success!" % (save_folder, filename))
except:
pass
if __name__ == '__main__':
# 1.参数设置
gray_thresh = 50 # 灰度图阈值
cut_wh = [2500, 1000] # 裁切的区域长宽
min_area, max_area = 100000, 2000000 # 最小、最大过滤面积
folder = r"C:\Users\HJ\Desktop\iabfkkcn" # 原图路径
save_folders = r"C:\Users\HJ\Desktop\归一化测试\output_image" # 保存路径
for root, dirs, files in os.walk(folder):
if root == folder:
pass
else:
sub_folder_name = os.path.join(save_folders, root.split("\\")[-1])
if not os.path.exists(sub_folder_name):
os.makedirs(sub_folder_name)
for file in files:
if file.split(".")[-1] in ["jpg", "png", "bmp"]:
origin_img_path = os.path.join(root, file)
image_normalize(origin_img_path, sub_folder_name, file, gray_thresh, cut_wh, min_area, max_area)
else:
continue