4nt-space / inference_utils.py
ki4n-4nt's picture
Removed redundant preprocessing
ce57502
Raw
History Blame Contribute Delete
2.37 kB
import tensorflow as tf
from tensorflow import keras
import numpy as np
import cv2
from PIL import Image
import json
# Constants
IMG_SIZE = 224
MODEL_PATH = "model/plant_disease_efficientnetb0.weights.h5"
CLASS_NAMES_PATH = "model/class_names.json"
# Load CLASS_NAMES
with open(CLASS_NAMES_PATH, "r") as f:
CLASS_NAMES = json.load(f)
def build_model(num_classes, img_size=IMG_SIZE):
inputs = keras.Input(shape=(img_size, img_size, 3))
base_model = keras.applications.EfficientNetB0(
include_top=False,
weights=None,
input_shape=(img_size, img_size, 3)
)
base_model.trainable = False
x = keras.applications.efficientnet.preprocess_input(inputs)
x = base_model(x, training=False)
x = keras.layers.GlobalAveragePooling2D()(x)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.Dropout(0.3)(x)
outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
return keras.Model(inputs, outputs)
NUM_CLASSES = len(CLASS_NAMES)
model = build_model(NUM_CLASSES)
model.load_weights(MODEL_PATH)
# Preprocess image
def preprocess_image(image_path):
img = tf.keras.preprocessing.image.load_img(
image_path, target_size=(IMG_SIZE, IMG_SIZE)
)
img_array = tf.keras.preprocessing.image.img_to_array(img)
# Preprocessing will be done by the model's input layer
return np.expand_dims(img_array, axis=0)
# Inference
def predict_plant_disease(image_path):
img_array = preprocess_image(image_path)
preds = model.predict(img_array)[0]
class_index = int(np.argmax(preds))
confidence = float(preds[class_index])
label = CLASS_NAMES[class_index]
return {label: confidence}
if __name__ == "__main__":
print("Model loaded. Enter image paths to classify (Ctrl+C to exit):\n")
try:
while True:
image_path = input("Enter image path: ").strip()
if not image_path:
print("Please enter a valid path.\n")
continue
try:
result = predict_plant_disease(image_path)
for label, confidence in result.items():
print(f"Label: {label}, Confidence: {confidence:.4f}\n")
except Exception as e:
print(f"Error processing image: {e}\n")
except KeyboardInterrupt:
print("\n\nExiting...")