Use the code below to get started with the model.
1
2
3import numpy as np
4import tensorflow as tf
5from tensorflow.keras.layers import *
6from tensorflow.keras.preprocessing.image import ImageDataGenerator
7from tensorflow.keras.models import Model
8from keras_self_attention import SeqSelfAttention
9from sklearn.metrics import *
10import seaborn as sns
11import matplotlib.pyplot as plt
12from tensorflow.keras.applications.efficientnet_v2 import EfficientNetV2B0
13
14
15input_shape = (sz, sz, 3) # Adjust the shape according to your data
16input_layer = Input(shape=input_shape)
17base_model = EfficientNetV2B0(include_top=False, weights='imagenet', input_tensor=input_layer).output
18x = GlobalAveragePooling2D()(base_model)
19x = Dense(64,activation = 'relu')(x)
20x = Dense(32,activation = 'relu')(x)
21x = Dropout(0.1)(x)
22num_classes = 6 # Six, including 'Other' label
23output_layer = Dense(num_classes, activation='softmax')(x)
24
25model = Model(inputs=input_layer, outputs=output_layer)
26model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
27model.load_weights("./EfficientNetV2B0_BestWeights.h5")
28
29data_dir = "./FetalClassification/Split_Images" # https://zenodo.org/records/3904280
30train_data_dir = f'{data_dir}/train'
31test_data_dir = f'{data_dir}/val'
32
33sz = 256
34bs = 16
35
36train_datagen = ImageDataGenerator(
37 rotation_range=30,
38 zoom_range=0.2,
39 horizontal_flip=True,
40 vertical_flip=True,
41)
42
43test_datagen = ImageDataGenerator()
44
45# Create train and test datasets using image_dataset_from_directory
46train_generator = train_datagen.flow_from_directory(
47 train_data_dir,
48 target_size=(sz,sz),
49 batch_size=bs,
50 class_mode = 'categorical',
51 seed=42, # Strictly setting seed for reproducibility
52)
53
54test_generator = test_datagen.flow_from_directory(
55 test_data_dir,
56 target_size=(sz,sz),
57 batch_size=bs,
58 class_mode = 'categorical',
59 seed=42, # Strictly setting seed for reproducibility
60 shuffle=False,
61)
62
63predictions = model.predict(test_generator)
64predicted_labels = np.argmax(predictions,axis = 1)
65true_labels = test_generator.classes
66
67classification_report_result = classification_report(true_labels, predicted_labels,digits = 4)
68confusion_matrix_result = confusion_matrix(true_labels, predicted_labels)
69
1@article{2410.17396,
2Author = {Arrun Sivasubramanian and Divya Sasidharan and Sowmya V and Vinayakumar Ravi},
3Title = {Efficient Feature Extraction Using Light-Weight CNN Attention-Based Deep Learning Architectures for Ultrasound Fetal Plane Classification},
4Year = {2024},
5Eprint = {arXiv:2410.17396},
6Doi = {10.1007/s13246-025-01566-6},
7}
