效果如图:
#文件夹预测
import math
from unet import Unet
from PIL import Image
import numpy
import os
import os.path
import numpy as np
import cv2
unet = Unet()
# 指明被遍历的文件夹
maindir = r'G:\\DL\\unet-keras\\unet-keras-master\\stone_predict\\ture\\stone'
blenddir= r'G:\\DL\\unet-keras\\unet-keras-master\\stone_predict\\ture\\add1207'
def drawDetect(img,predict_image,image_bgr):
shape = predict_image.shape
w = shape[1]
h = shape[0]
threshold=h/30 * w/30
blurred = cv2.GaussianBlur(img, (11, 11), 0)
edged = cv2.Canny(blurred, 30, 150) # 用Canny算子提取边缘
contours, hierarchy = cv2.findContours(edged.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # 寻找封闭区域
maxCnt=-1
if contours is not None:
for i in range(len(contours)):
area = cv2.contourArea(contours[i]) # 计算轮廓所占面积
if area > threshold: # 将area小于阈值区域填充背景色,由于OpenCV读出的是BGR值
# cv2.drawContours(img, [contours[i]], -1, (84, 1, 68), thickness=-1) # 原始图片背景BGR值(84,1,68)
rect = cv2.minAreaRect(contours[i]) # 最小外接矩形
rect_w, rect_h = (rect[1][0]) + 1.0, (rect[1][1]) + 1.0 # 最小外接矩形的宽和高
if rect_w > rect_h:
angle = 90 - rect[2]
L = rect_w
Ly = rect_w * math.cos(math.radians(angle))
Lx = rect_w * math.sin(math.radians(angle))
else:
angle = rect[2]
L = rect_h
Ly = rect_h * math.cos(math.radians(angle))
Lx = rect_h * math.sin(math.radians(angle))
box = np.int0(cv2.boxPoints(rect)) # 矩形的四个角点取整
if box is not None:
image_bgr = cv2.drawContours(image_bgr, [box], 0, (0, 0, 255), 2)
image_bgr = cv2.drawContours(image_bgr, contours[i], -1, (0, 255, 0), 3)
continue
# predict_image = cv2.drawContours(image_bgr, contours, -1, (0, 255, 0), 3)
# return contours, box
return image_bgr
# while True:
for parent, dirnames, filenames in os.walk(maindir): # 遍历每一张图片
for filename in filenames:
pic_name = os.path.join(parent, filename)
image = Image.open(pic_name)
r_image = unet.detect_image(image)
# predict_image = cv2.imread(r_image)
image_bgr = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR) # PIL转为cv.bgr
predict_image = cv2.cvtColor(numpy.asarray(r_image), cv2.COLOR_RGB2BGR) # PIL转为cv.bgr
predict_gray = cv2.cvtColor(predict_image, cv2.COLOR_BGR2GRAY) # 转为灰度图
# contours, box = drawDetect(predict_gray,predict_image)
predict_image2 = drawDetect(predict_gray,predict_image,image_bgr)
if predict_image2 is None:
predict_image2 = image_bgr
# print(type(predict_image2))
blend_image = Image.fromarray(cv2.cvtColor(predict_image2,cv2.COLOR_BGR2RGB))
blend_image.save(blenddir +'\\' + pic_name.split('.')[0].strip(parent) + ".png") # 存储裁剪得到的图像,需要先创建文件夹