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