|
|
import os |
|
|
import urllib.request as request |
|
|
from zipfile import ZipFile |
|
|
import tensorflow as tf |
|
|
from cnnClassfier.entity.config_entity import PrepareBaseModelConfig |
|
|
from pathlib import Path |
|
|
|
|
|
|
|
|
|
|
|
class PrepareBaseModel: |
|
|
def __init__(self, config: PrepareBaseModelConfig): |
|
|
self.config = config |
|
|
|
|
|
def get_base_model(self): |
|
|
self.model = tf.keras.applications.vgg16.VGG16( |
|
|
input_shape = self.config.params_image_size, |
|
|
weights = self.config.params_weights, |
|
|
include_top = self.config.params_include_top |
|
|
) |
|
|
|
|
|
self.save_model(path = self.config.base_model_path, model = self.model) |
|
|
|
|
|
|
|
|
@staticmethod |
|
|
def prepare_full_model(model, classes, freeze_all, freeze_till, learinig_rate): |
|
|
if freeze_all: |
|
|
for layer in model.layers: |
|
|
model.trainable = False |
|
|
|
|
|
elif (freeze_till is not None) and (freeze_till > 0): |
|
|
for layer in model.layers[:-freeze_till]: |
|
|
model.trainable = False |
|
|
|
|
|
flatten_in = tf.keras.layers.Flatten()(model.output) |
|
|
prediction = tf.keras.layers.Dense( |
|
|
units = classes, |
|
|
activation = 'softmax' |
|
|
)(flatten_in) |
|
|
|
|
|
full_model = tf.keras.models.Model( |
|
|
inputs = model.input, |
|
|
outputs = prediction |
|
|
) |
|
|
|
|
|
full_model.compile( |
|
|
optimizer = tf.keras.optimizers.SGD(lr = learinig_rate), |
|
|
loss = tf.keras.losses.CategoricalCrossentropy(), |
|
|
metrics = ['accuracy'] |
|
|
) |
|
|
|
|
|
full_model.summary() |
|
|
return full_model |
|
|
|
|
|
def update_base_model(self): |
|
|
self.full_model = self.prepare_full_model( |
|
|
model = self.model, |
|
|
classes = self.config.params_classes, |
|
|
freeze_all=True, |
|
|
freeze_till=None, |
|
|
learinig_rate=self.config.params_learning_rate |
|
|
) |
|
|
self.save_model(path = self.config.updated_base_model_path, model = self.full_model) |
|
|
|
|
|
|
|
|
@staticmethod |
|
|
def save_model(path: Path, model: tf.keras.Model): |
|
|
model.save(path) |
|
|
|
