|
- 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")
|