python opencv提取图中矩形框及其位置并保存到json中

矩形框检测流程：

1. 图像转为灰度图
2. canny边缘检测
3. 霍夫变换提取直线
4. 多边形拟合找到矩形框
5. 计算矩形框的最小外接矩形作为新的标注框
6. 写入json

import base64
import cv2
import numpy as np
import json
import os
from skimage import data,color,morphology

img_path = './Data/images'
save_json_dir = './Data/json'
label_name = 'face'


def save_jsons(imagePath,bbox,imageHeight,imageWidth):
    cur_json_dict = {
        "version": "4.5.7",
        "flags": {},
        "shapes": [],
    }

    # bbox
    for box in bbox:
        box = [[float(box[0][0]),float(box[0][1])],[float(box[2][0]),float(box[2][1])]]
        cur_json_dict['shapes'].append(
                {"label": label_name, 
                 "points": box, 
                 'group_id': None, 
                 "shape_type": "rectangle", 
                 "flags": {}
                 })

    cur_json_dict["imagePath"] = imagePath
    cur_json_dict["imageData"] = None
    cur_json_dict["imageHeight"] = imageHeight
    cur_json_dict["imageWidth"] = imageWidth
    with open(save_json_dir+'/'+os.path.basename(imagePath)[:-4]+'.json', 'w') as f:
        f.write(json.dumps(cur_json_dict))


def find_rec(img_path):
    img_list = os.listdir(img_path)
    for img_name in img_list:
        save_p = './saveTemp' + '/' + img_name
        img_p = os.path.join(img_path,img_name)
        print(img_p)
        img_rgb = cv2.imread(img_p)
        img_h,img_w,img_c = img_rgb.shape

        # 1. 图像转为灰度图
        img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
        # cv2.imshow('img',img_gray)
        # cv2.waitKey()

        # 2. canny边缘检测 
        edges = cv2.Canny(img_gray,200,300)
        # cv2.imwrite('./1.jpg',edges)
        # cv2.imwrite(save_p,edges)

        # 3. 霍夫变换提取直线，并用白线画到和原图大小一致的黑底图上
        # # lines = cv2.HoughLines(edges,rho=1,theta=np.pi/180,threshold=118)
        linesv = cv2.HoughLinesP(edges, rho=1,theta=np.pi,threshold=25,minLineLength=25,maxLineGap=10)
        linesh = cv2.HoughLinesP(edges, rho=1,theta=np.pi/2,threshold=25,minLineLength=25,maxLineGap=10)
        lines =  np.append(linesv,linesh,axis=0)
        
        # 绘制纯黑图
        img_black = np.zeros(img_rgb.shape, dtype=np.uint8)
        # cv2.imwrite(save_p,img_black)
        
        # 在黑底图上用白色画提取到的直线
        for l in lines.tolist():
                x1, y1, x2, y2 = l[0]
                cv2.line(img_black, (x1, y1), (x2, y2), color=(255, 255, 255), thickness=2)
        # cv2.imwrite(save_p,img_black)

        black_gray = cv2.cvtColor(img_black, cv2.COLOR_RGB2GRAY)

        # 找出图中的轮廓值
        contours, hierarchy = cv2.findContours(black_gray,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        bbox = []
        for contour in contours:
            # 4. 多边形拟合找到矩形框
            approx1 = cv2.approxPolyDP(contour,10,True) # 轮廓线的多边形逼近 中间的epsilon是拟合精确度 最后是True指定是否是闭合曲线
            # img_black  = cv2.polylines(img_black,[approx1],True,(0,255,0),2)

            # 5. 计算矩形框的最小外接矩形作为新的标注框
            rect = cv2.minAreaRect(approx1)
            box = cv2.boxPoints(rect) # 得到最小矩形的坐标
            box = np.int0(box) # 标准化坐标到整数
            # print(box)
            # cv2.drawContours(img_black, [box], 0, (0, 255, 0), 3)
            
            # 去掉错把圆的边缘检测为直线的部分
            if abs(box[0][0] - box[2][0]) > 10 and abs(box[0][1] - box[2][1]) > 10:
                # x1,y1,x2,y2 = box[0][0],box[0][1],box[2][0],box[2][1]
                x1 = min(box[0][0], box[2][0])
                x2 = max(box[0][0], box[2][0])
                y1 = min(box[0][1], box[2][1])
                y2 = max(box[0][1], box[2][1])
                cv2.rectangle(img_black,(x1,y1),(x2,y2),(0,255,0),2)
                bbox.append(box)

        # cv2.imwrite(save_p,img_black)

        # 6. 写入json
        save_jsons(img_p,bbox,img_h,img_w)


if __name__ == '__main__':
    find_rec(img_path)