code_test/AOI_select.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')
    args = parser.parse_args()
    return args
    
    
args = parse_args()
#path = "D:/Work/20211119-dsxw/PCBA_dataset/dsxw_report/SMD12-09-Report/SMD12-Manual/img/ng/"
path = args.AOI_path+"/ng/"
print(path)
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"}
class_name_other = {"上锡不良_Poor_S":"12","qita":"13","limit_error":"13","极性错_Wrong_Po":"13","其它_Others":"13"}
count = 0
count_qita = 0
bpp_ok = 0
bpp_ng = 0
count_ng = 0
def cv_imread(file_path):
    cv_img = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), 0)
    return cv_img
    
for img in imgs:
    i = i+1
    src = cv_imread(path+img)
    #print(img)
    img_name = img.split("@")
    #print(img_name[2], img_name[3])
    if img_name[2] in class_AOI_name.keys():
        #print(img_name[2], img)
        if img_name[3]=="Pseudo_Error":
            bpp_ok += 1
        else:
            bpp_ng += 1
        count += 1
        os.remove(path+img)
        continue
    if img_name[3] not in class_name.keys():
        print(img_name[3], img)
        count_qita += 1
        os.remove(path+img)
        continue
    count_ng += 1
print(count, bpp_ok, bpp_ng)
print(count_qita)
print(count_ng)