import cv2
import numpy as np
import os
def get_bbox(file_path):
for subdir in os.listdir(file_path):
path=file_path+ '/'+subdir
mask_path=file_path+'_label/'+subdir
if not os.path.exists(mask_path):
os.makedirs(mask_path)
for name in os.listdir(path):
img_p=path+'/'+name
bgr_img = cv2.imread(img_p)
lower_blue=np.array([10,10,10])
upper_blue=np.array([255,255,255])
mask = cv2.inRange(bgr_img, lower_blue, upper_blue)
kernel = np.ones((5,5),np.uint8)
mask = cv2.dilate(mask, kernel, iterations = 1) #膨脹操作
contours, hierarchy = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) #找出主體的輪廓
# 生成mask
img_gray = cv2.cvtColor(bgr_img,cv2.COLOR_BGR2GRAY)
mask = np.ones_like(img_gray)
area = []
for k in range(len(contours)):
area.append(cv2.contourArea(contours[k]))
max_idx = np.argmax(np.array(area))
cv2.drawContours(mask, contours, max_idx, 255, cv2.FILLED)
cv2.imwrite(mask_path+'/'+name, mask)
if __name__ == "__main__":
img_path='E:/datasets/yanye_shixiong/classify/val_org'
get_bbox(img_path)
得到的結果:
(image)
(得到的mask)
(以一張圖像爲例)