digit_classifier.py 文件源码

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)