Finetuned-MobilVIT / Train_code_mobilenet
Nekshay's picture
Update Train_code_mobilenet
588994d
raw
history blame
2.8 kB
import tensorflow as tf
from tensorflow.keras.applications import MobileNet
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# Load the MobileNet base model
base_model = MobileNet(weights='imagenet', include_top=False)
# Add custom classification layers
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x)
num_classes=2
predictions = Dense(num_classes, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=predictions)
# Compile the model
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
# Data augmentation and preprocessing
train_datagen = ImageDataGenerator(
preprocessing_function=tf.keras.applications.mobilenet.preprocess_input,
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
horizontal_flip=True
)
batch_size=16
train_generator = train_datagen.flow_from_directory(
'/content/tire-dataset/train_data',
target_size=(224, 224),
batch_size=batch_size,
class_mode='categorical'
)
test_datagen = ImageDataGenerator(
preprocessing_function=tf.keras.applications.mobilenet.preprocess_input,
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
horizontal_flip=True
)
batch_size=16
# Train the model
num_epochs=1
model.fit(train_generator, epochs=num_epochs)
# Evaluate the model on the test set
test_generator = test_datagen.flow_from_directory(
'/content/tire-dataset/test_data',
target_size=(224, 224),
batch_size=batch_size,
class_mode='categorical'
)
accuracy = model.evaluate(test_generator)
print('Test accuracy:', accuracy)
from tensorflow import keras
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictions
import numpy as np
# Load the model
#model = keras.models.load_model('path_to_your_model.h5')
# Load and preprocess an image for inference
img_path = '/content/tire-dataset/test_data/Tire/00000.jpg'
img = image.load_img(img_path, target_size=(224, 224))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)
# Make a prediction
predictions = model.predict(x)
# Decode and display the prediction
# decoded_predictions = decode_predictions(predictions, top=3)[0]
# for label, description, score in decoded_predictions:
# print(f'{label}: {description} ({score:.2f})')
model.save('/content/model_keras/keras_model.h5')
!tensorflowjs_converter --input_format=keras --output_format=tfjs_graph_model --split_weights_by_layer --weight_shard_size_bytes=99999999 --quantize_float16=* /content/model_keras/keras_model.h5 ./model_tfjs