Spaces:
Running
Running
| import keras | |
| from keras.applications.resnet50 import ResNet50 | |
| from keras.models import Model,load_model | |
| from keras.layers import Dense, Activation, Flatten, Dropout | |
| from sklearn.metrics import confusion_matrix | |
| import numpy as np | |
| import cv2 | |
| from config.CONFIG import Config | |
| cfg = Config() | |
| class CNN: | |
| def __init__(self, classifier_type, is_load=True): | |
| ''' | |
| :param classifier_type: 'Text' or 'Noise' or 'Elements' | |
| ''' | |
| self.data = None | |
| self.model = None | |
| self.classifier_type = classifier_type | |
| self.image_shape = (32,32,3) | |
| self.class_number = None | |
| self.class_map = None | |
| self.model_path = None | |
| self.classifier_type = classifier_type | |
| if is_load: | |
| self.load(classifier_type) | |
| def build_model(self, epoch_num, is_compile=True): | |
| base_model = ResNet50(include_top=False, weights='imagenet', input_shape=self.image_shape) | |
| for layer in base_model.layers: | |
| layer.trainable = False | |
| self.model = Flatten()(base_model.output) | |
| self.model = Dense(128, activation='relu')(self.model) | |
| self.model = Dropout(0.5)(self.model) | |
| self.model = Dense(15, activation='softmax')(self.model) | |
| self.model = Model(inputs=base_model.input, outputs=self.model) | |
| if is_compile: | |
| self.model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy']) | |
| self.model.fit(self.data.X_train, self.data.Y_train, batch_size=64, epochs=epoch_num, verbose=1, | |
| validation_data=(self.data.X_test, self.data.Y_test)) | |
| def train(self, data, epoch_num=30): | |
| self.data = data | |
| self.build_model(epoch_num) | |
| self.model.save(self.model_path) | |
| print("Trained model is saved to", self.model_path) | |
| def load(self, classifier_type): | |
| if classifier_type == 'Text': | |
| self.model_path = 'E:/Mulong/Model/rico_compos/cnn-textview-2.h5' | |
| self.class_map = ['Text', 'Non-Text'] | |
| elif classifier_type == 'Noise': | |
| self.model_path = 'E:/Mulong/Model/rico_compos/cnn-noise-1.h5' | |
| self.class_map = ['Noise', 'Non-Noise'] | |
| elif classifier_type == 'Elements': | |
| # self.model_path = 'E:/Mulong/Model/rico_compos/resnet-ele14-19.h5' | |
| # self.model_path = 'E:/Mulong/Model/rico_compos/resnet-ele14-28.h5' | |
| # self.model_path = 'E:/Mulong/Model/rico_compos/resnet-ele14-45.h5' | |
| self.model_path = 'UIED/cnn/model/cnn-rico-1.h5' # Use local model | |
| self.class_map = cfg.element_class | |
| self.image_shape = (64, 64, 3) | |
| elif classifier_type == 'Image': | |
| # Redirect 'Image' classification to use the general 'Elements' model | |
| # as the specific model is not available in the project. | |
| # IMPORTANT: This requires the actual model file to be present for real classification. | |
| print("Warning: 'Image' specific model not found. Redirecting to general 'Elements' classifier.") | |
| self.model_path = 'UIED/cnn/model/cnn-rico-1.h5' # Use local model | |
| self.class_map = ['Image', 'Non-Image'] # Keep the class map for binary classification logic | |
| self.class_number = len(self.class_map) | |
| try: | |
| self.model = load_model(self.model_path) | |
| print('Model Loaded From', self.model_path) | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| print("A dummy model file was created, but it's not a valid Keras model.") | |
| print("Please replace it with the actual model file for classification to work.") | |
| self.model = None | |
| def preprocess_img(self, image): | |
| image = cv2.resize(image, self.image_shape[:2]) | |
| x = (image / 255).astype('float32') | |
| x = np.array([x]) | |
| return x | |
| def predict(self, imgs, compos, load=False, show=False): | |
| """ | |
| :type img_path: list of img path | |
| """ | |
| if load: | |
| self.load(self.classifier_type) | |
| if self.model is None: | |
| print("*** No model loaded ***") | |
| return | |
| for i in range(len(imgs)): | |
| X = self.preprocess_img(imgs[i]) | |
| Y = self.class_map[np.argmax(self.model.predict(X))] | |
| compos[i].category = Y | |
| if show: | |
| print(Y) | |
| cv2.imshow('element', imgs[i]) | |
| cv2.waitKey() | |
| def evaluate(self, data, load=True): | |
| if load: | |
| self.load(self.classifier_type) | |
| X_test = data.X_test | |
| Y_test = [np.argmax(y) for y in data.Y_test] | |
| Y_pre = [np.argmax(y_pre) for y_pre in self.model.predict(X_test, verbose=1)] | |
| matrix = confusion_matrix(Y_test, Y_pre) | |
| print(matrix) | |
| TP, FP, FN = 0, 0, 0 | |
| for i in range(len(matrix)): | |
| TP += matrix[i][i] | |
| FP += sum(matrix[i][:]) - matrix[i][i] | |
| FN += sum(matrix[:][i]) - matrix[i][i] | |
| precision = TP/(TP+FP) | |
| recall = TP / (TP+FN) | |
| print("Precision:%.3f, Recall:%.3f" % (precision, recall)) |