| import cv2 | |
| import numpy as np | |
| import tensorflow as tf | |
| from tensorflow.keras.optimizers import Adam | |
| from tensorflow.keras.models import Sequential | |
| from tensorflow.keras.layers import Dense, Flatten | |
| class RecycleModel: | |
| def __init__(self): | |
| self.vgg16_model = Sequential() | |
| self.pretrained_model = tf.keras.applications.VGG16( | |
| include_top=False, | |
| weights="imagenet", | |
| input_shape=(180, 180, 3), | |
| pooling='max', | |
| classes=7 | |
| ) | |
| for layer in self.pretrained_model.layers: | |
| layer.trainable = False | |
| self.vgg16_model.add(self.pretrained_model) | |
| self.vgg16_model.add(Flatten()) | |
| self.vgg16_model.add(Dense(512, activation='relu')) | |
| self.vgg16_model.add(Dense(7, activation='softmax')) | |
| self.vgg16_model.compile(optimizer=Adam(learning_rate=0.001), | |
| loss='sparse_categorical_crossentropy', | |
| metrics=['accuracy']) | |
| self.vgg16_model.load_weights('updated_vgg16_weights.h5') | |
| self.class_names = ['1_polyethylene_PET', '2_high_density_polyethylene_PE-HD', '3_polyvinylchloride_PVC', | |
| '4_low_density_polyethylene_PE-LD', '5_polypropylene_PP', '6_polystyrene_PS', '8_no_plastic'] | |
| def recyclePrediction(self, img): | |
| image_resized = cv2.resize(img, (180, 180)) | |
| image = np.expand_dims(image_resized, axis=0) | |
| pred = self.vgg16_model.predict(image) | |
| output_class = self.class_names[np.argmax(pred)] | |
| return output_class | |