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