Spaces:
Sleeping
Sleeping
File size: 1,363 Bytes
47488ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import tensorflow as tf
from pathlib import Path
from kidney_classification.entity.config_entity import PrepareBaseModelConfig
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, BatchNormalization, Dropout
class PrepareBaseModel:
@staticmethod
def prepare_full_model():
VGG_model = Sequential()
pretrained_model = tf.keras.applications.VGG16(
include_top=False,
input_shape=(150, 150, 3),
pooling="max",
classes=4,
weights="imagenet",
)
VGG_model.add(pretrained_model)
VGG_model.add(Flatten())
VGG_model.add(Dense(512, activation="relu"))
VGG_model.add(BatchNormalization())
VGG_model.add(Dropout(0.5))
VGG_model.add(Dense(4, activation="softmax"))
pretrained_model.trainable = False
VGG_model.compile(
optimizer="adam",
loss="sparse_categorical_crossentropy",
metrics=["accuracy"],
)
return VGG_model
def update_base_model(self, config: PrepareBaseModelConfig):
full_model = self.prepare_full_model()
full_model.summary()
full_model.save(config.updated_base_model_path)
@staticmethod
def save_model(path: Path, model: tf.keras.Model):
model.save(path)
|