rice-scanner / task_2_model.py
NickNam2710's picture
update load model
1413913
import os
import tensorflow as tf
from huggingface_hub import hf_hub_download
@tf.keras.utils.register_keras_serializable()
class InceptionModel(tf.keras.Model):
def __init__(
self,
dropout_rate: float,
l2_reg: float,
dense_units: int,
*,
name="InceptionV3",
**kwargs,
):
super().__init__(name=name, **kwargs)
# Store for get_config
self.dropout_rate = dropout_rate
self.l2_reg = l2_reg
self.dense_units = dense_units
# Ensure L2 regularizer is correctly instantiated
l2_regularizer = tf.keras.regularizers.L2(l2_reg)
inception_base = tf.keras.applications.InceptionV3(
include_top=False, # Keep False to add custom top layers
input_shape=(256, 256, 3),
pooling='max',
weights=None
)
inputs = tf.keras.layers.Input(shape=(256, 256, 3))
x = tf.keras.layers.Rescaling(1./255.)(inputs) # Scale pixel values to [0, 1]
x = inception_base(x) # Pass through the InceptionV3 base
x = tf.keras.layers.Dropout(dropout_rate)(x)
x = tf.keras.layers.Dense(dense_units, activation="relu", kernel_regularizer=l2_regularizer)(x)
outputs = tf.keras.layers.Dense(10, activation='softmax')(x) # Output probabilities for each class
# Create the Keras Model
self.net = tf.keras.Model(inputs=inputs, outputs=outputs)
def call(self, inputs, training=False):
return self.net(inputs, training=training)
def get_config(self):
config = super().get_config()
config.update({
"dropout_rate": self.dropout_rate,
"l2_reg": self.l2_reg,
"dense_units": self.dense_units,
})
return config
@classmethod
def from_config(cls, config):
dropout_rate = config.pop("dropout_rate")
l2_reg = config.pop("l2_reg")
dense_units = config.pop("dense_units")
return cls(dropout_rate, l2_reg, dense_units, **config)
def load_task_2_model(model_name='task_2_model_inception.keras'):
model_path = hf_hub_download(
repo_id="NickNam2710/predict_rice_diseases",
filename=model_name,
revision="main"
)
model = tf.keras.models.load_model(model_path)
return model