Python cv2 模块，KNearest() 实例源码

def __init__(self,
                 saved_model=None,
                 train_folder=None,
                 feature=_feature.__func__):
        """
        :param saved_model: optional saved train set and labels as .npz
        :param train_folder: optional custom train data to process
        :param feature: feature function - compatible with saved_model
        """
        self.feature = feature
        if train_folder is not None:
            self.train_set, self.train_labels, self.model = \
                self.create_model(train_folder)
        else:
            if cv2.__version__[0] == '2':
                self.model = cv2.KNearest()
            else:
                self.model = cv2.ml.KNearest_create()
            if saved_model is None:
                saved_model = TRAIN_DATA+'raw_pixel_data.npz'
            with np.load(saved_model) as data:
                self.train_set = data['train_set']
                self.train_labels = data['train_labels']
                if cv2.__version__[0] == '2':
                    self.model.train(self.train_set, self.train_labels)
                else:
                    self.model.train(self.train_set, cv2.ml.ROW_SAMPLE,
                                     self.train_labels)

def create_model(self, train_folder):
        """
        Return the training set, its labels and the trained model
        :param train_folder: folder where to retrieve data
        :return: (train_set, train_labels, trained_model)
        """
        digits = []
        labels = []
        for n in range(1, 10):
            folder = train_folder + str(n)
            samples = [pic for pic in os.listdir(folder)
                       if os.path.isfile(os.path.join(folder, pic))]

            for sample in samples:
                image = cv2.imread(os.path.join(folder, sample))
                # Expecting black on white
                image = 255 - cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                _, image = cv2.threshold(image, 0, 255,
                                         cv2.THRESH_BINARY + cv2.THRESH_OTSU)
                feat = self.feature(image)
                digits.append(feat)
                labels.append(n)

        digits = np.array(digits, np.float32)
        labels = np.array(labels, np.float32)
        if cv2.__version__[0] == '2':
            model = cv2.KNearest()
            model.train(digits, labels)
        else:
            model = cv2.ml.KNearest_create()
            model.train(digits, cv2.ml.ROW_SAMPLE, labels)
        return digits, labels, model

def __init__(self):
        collect_dir = 'captcha/collect'
        label = []
        train_file = []
        for i in os.listdir(collect_dir):
            for y in os.listdir(collect_dir + '/' + i):
                #print i
                label.append(ord(i))
                #print y
                train_file.append(collect_dir + '/' + i + '/' + y)
        train_data = [cv2.imread(i, 0) for i in train_file]
        train = np.array(train_data).reshape(-1, 400).astype(np.float32)
        label = np.array(label).reshape(-1)
        self.knn = cv2.KNearest()
        self.knn.train(train, label)

def get_matrix():
    chessboard_matrix=init_matrix()
    train_filename_list,train_label_list=train.find_picture('static/train/')
    train_file_list = train.preprocess_img(train_filename_list)

    test_filename_list,test_label_list=train.find_picture('static/ClippedImg/')
    test_file_list = train.preprocess_img(test_filename_list)

    knn = cv2.KNearest()
    knn.train(train_file_list,train_label_list)
    ret,result,neighbours,dist = knn.find_nearest(test_file_list,k=3)
    for i in  range(len(result)):
        #print int(result[i][0]),test_filename_list[i]
        if int(result[i][0]) ==1:
            position = re.search(r'\d{1,2}_\d{1,2}',str(test_filename_list[i])).group().split('_')
            xposition,yposition = int(position[0]),int(position[1])
            chessboard_matrix[xposition][yposition]=1
        if int(result[i][0]) ==2:
            position = re.search(r'\d{1,2}_\d{1,2}',str(test_filename_list[i])).group().split('_')
            xposition,yposition = int(position[0]),int(position[1])
            chessboard_matrix[xposition][yposition]=2
        # if str(test_filename_list[i]).find('2_2') != -1:
        #     print result[i],test_filename_list[i]
    return chessboard_matrix