Efficient Feature Extraction Using Light-Weight CNN Attention-Based Deep Learning Architectures for Ultrasound Fetal Plane Classification
Paper โข 2410.17396 โข Published
Use the code below to get started with the model.
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Model
from keras_self_attention import SeqSelfAttention
from sklearn.metrics import *
import seaborn as sns
import matplotlib.pyplot as plt
from tensorflow.keras.applications.efficientnet_v2 import EfficientNetV2B0
input_shape = (sz, sz, 3) # Adjust the shape according to your data
input_layer = Input(shape=input_shape)
base_model = EfficientNetV2B0(include_top=False, weights='imagenet', input_tensor=input_layer).output
x = GlobalAveragePooling2D()(base_model)
x = Dense(64,activation = 'relu')(x)
x = Dense(32,activation = 'relu')(x)
x = Dropout(0.1)(x)
num_classes = 6 # Six, including 'Other' label
output_layer = Dense(num_classes, activation='softmax')(x)
model = Model(inputs=input_layer, outputs=output_layer)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.load_weights("./EfficientNetV2B0_BestWeights.h5")
data_dir = "./FetalClassification/Split_Images" # https://zenodo.org/records/3904280
train_data_dir = f'{data_dir}/train'
test_data_dir = f'{data_dir}/val'
sz = 256
bs = 16
train_datagen = ImageDataGenerator(
rotation_range=30,
zoom_range=0.2,
horizontal_flip=True,
vertical_flip=True,
)
test_datagen = ImageDataGenerator()
# Create train and test datasets using image_dataset_from_directory
train_generator = train_datagen.flow_from_directory(
train_data_dir,
target_size=(sz,sz),
batch_size=bs,
class_mode = 'categorical',
seed=42, # Strictly setting seed for reproducibility
)
test_generator = test_datagen.flow_from_directory(
test_data_dir,
target_size=(sz,sz),
batch_size=bs,
class_mode = 'categorical',
seed=42, # Strictly setting seed for reproducibility
shuffle=False,
)
predictions = model.predict(test_generator)
predicted_labels = np.argmax(predictions,axis = 1)
true_labels = test_generator.classes
classification_report_result = classification_report(true_labels, predicted_labels,digits = 4)
confusion_matrix_result = confusion_matrix(true_labels, predicted_labels)
@article{2410.17396,
Author = {Arrun Sivasubramanian and Divya Sasidharan and Sowmya V and Vinayakumar Ravi},
Title = {Efficient Feature Extraction Using Light-Weight CNN Attention-Based Deep Learning Architectures for Ultrasound Fetal Plane Classification},
Year = {2024},
Eprint = {arXiv:2410.17396},
Doi = {10.1007/s13246-025-01566-6},
}
Base model
timm/tf_efficientnetv2_b0.in1k