code_test/AOI_get_box.py

import cv2
import os
import numpy as np
import shutil
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description='get AOI class')
    parser.add_argument('--AOI_path', help='AOI path')
    parser.add_argument('--coco_path', help='coco path')
    parser.add_argument('--classes_file', help='classes file path')
    args = parser.parse_args()
    return args

def setDir(filepath):
    '''
    如果文件夹不存在就创建，如果文件存在就清空！
    :param filepath:需要创建的文件夹路径
    :return:
    '''
    if not os.path.exists(filepath):
        os.mkdir(filepath)
    else:
        shutil.rmtree(filepath,ignore_errors=True)
        # os.mkdir(filepath)

args = parse_args()

path = args.AOI_path+"/ng/"
coco_path = args.coco_path
coco_ann_path = args.coco_path+"/annotations/"
coco_img_path = args.coco_path+"/images/"
labels_path = args.coco_path+"/labels/"
classes_file = args.classes_file
setDir(coco_path)
setDir(coco_ann_path)
setDir(coco_img_path)
setDir(labels_path)

def order_points(pts):
    # pts为轮廓坐标
    # 列表中存储元素分别为左上角，右上角，右下角和左下角
    rect = np.zeros((4, 2), dtype = "float32")
    # 左上角的点具有最小的和，而右下角的点具有最大的和
    s = pts.sum(axis = 1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    # 计算点之间的差值
    # 右上角的点具有最小的差值,
    # 左下角的点具有最大的差值
    diff = np.diff(pts, axis = 1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    # 返回排序坐标(依次为左上右上右下左下)
    return rect

    
def angle_cos(p0, p1, p2):
    d1, d2 = (p0-p1).astype('float'), (p2-p1).astype('float')
    return abs( np.dot(d1, d2) / np.sqrt( np.dot(d1, d1)*np.dot(d2, d2) ) )
    
def cv_rename(file_path = ""):
    file_path_gbk = file_path.encode('gbk')
    return file_path_gbk.decode()

imgs = os.listdir(path)
i = 0
class_AOI_name = {"bu_pi_pei":"1","fang_xiang_fan":"2","err.txt_c_not_f":"3"}
class_name = {"yi_wei":"1","lou_jian":"2","ce_li":"3","li_bei":"4","shang_xia_fan_t":"5","lian_xi":"6","duo_jian":"7","sun_huai":"8","shao_xi":"9","jia_han":"10","yi_wu":"11",\
              "移位_Component_":"1","缺件_Component_":"2","侧立_Stand_Up":"3","立碑_Tombstone":"4","翻贴_Upside_Dow":"5","连锡_Solder_Bri":"6","Solderbridge":"6",\
              "损坏_Bad_Compon":"8","少锡_Insufficie":"9","假焊_Pseudo_Sol":"10", "qi_ta": "10"}
count = 0
count_qita = 0
for img in imgs:
    i = i+1
    #print(img)
    img_name = img.split("@")
    #print(img_name[2], img_name[3])
    src = cv2.imdecode(np.fromfile(path+img, dtype=np.uint8), cv2.IMREAD_COLOR)# 默认的彩色图(IMREAD_COLOR)方式读入原始图像
    mask = cv2.imdecode(np.fromfile(path+img, dtype=np.uint8), 0)    # 灰度图(IMREAD_GRAYSCALE)方式读入水印蒙版图像
    mask = np.where(mask<240, 239, mask)
    mask = (np.divide(mask-239, 16)*255)
    mask = np.where(mask<60, 0, mask)
    mask = np.uint8(mask)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(3,3)) 
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel,iterations=3)
    
    contours, hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    #print(contours)
    squares = []
    result = []
    k = 0
    for cnt in contours:
        cnt_len = cv2.arcLength(cnt, True) #计算轮廓周长
        #cnt = cv2.approxPolyDP(cnt, 0.02*cnt_len, True) #多边形逼近
        # 条件判断逼近边的数量是否为4，轮廓面积是否大于1000，检测轮廓是否为凸的
        #print(len(cnt), cv2.contourArea(cnt), cv2.isContourConvex(cnt))
        if len(cnt) >= 4 and cv2.contourArea(cnt) > 30:
            #print("***********************")
            M = cv2.moments(cnt) #计算轮廓的矩
            cx = int(M['m10']/M['m00'])
            cy = int(M['m01']/M['m00'])#轮廓重心
            
            cnt = cnt.reshape(-1, 2)
            max_cos = np.max([angle_cos( cnt[i], cnt[(i+1) % 4], cnt[(i+2) % 4] ) for i in range(4)])
            # 只检测矩形（cos90° = 0）
            #if max_cos < 0.3:
            # 检测四边形（不限定角度范围）
            if True:
                #index = index + 1
                #cv2.putText(img,("#%d"%index),(cx,cy),font,0.7,(255,0,255),2)
                if cv2.contourArea(cnt)>k:
                    k = cv2.contourArea(cnt)
                    result = cnt
                squares.append(cnt)
    mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
    try:
        rect = order_points(result)
    except:
        print(img)
        os.remove(path+img)
        continue
    #print(rect)  # 获取最大内接矩形的4个顶点坐标
    x1 = 1000
    y1 = 1000
    x2 = 0
    y2 = 0
    for xy in rect:
        if x1>xy[0]:
            x1 = xy[0]
        if y1>xy[1]:
            y1 = xy[1]
        if x2<xy[0]:
            x2 = xy[0]
        if y2<xy[1]:
            y2 = xy[1]
    cv2.rectangle(src, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 0), 3)
    label_result = img+" "+str(int(x1))+","+str(int(y1))+","+str(int(x2))+","+str(int(y2))+","+class_name[img_name[3]]
    #shutil.copy(path+img, coco_path+img)
    #print(img, img[:-4])
    path_file_name = labels_path+img[:-4]+".txt"
    if not os.path.exists(path_file_name):
        with open(path_file_name, "a") as f:
            f.write(label_result)
    
    #break
print(class_name)
print(count, count_qita)
print("copying img")
ok_path = args.AOI_path+"/ok/"
ng_path = args.AOI_path+"/ng/"
imgs = os.listdir(ok_path)
for img in imgs:
    shutil.copy(ok_path+img, coco_img_path+img)

imgs = os.listdir(ng_path)
for img in imgs:
    shutil.copy(ng_path+img, coco_img_path+img)

shutil.copy(classes_file, coco_path+"classes.txt")