260 lines
11 KiB
Python
260 lines
11 KiB
Python
import os, tf2onnx, numpy, random, math
|
|
import tensorflow as tf
|
|
|
|
from tensorflow.keras import *
|
|
from PIL import Image
|
|
|
|
####################################################################################################
|
|
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
tf.keras.mixed_precision.set_global_policy('mixed_float16')
|
|
|
|
####################################################################################################
|
|
|
|
def get_file_list(directory):
|
|
file_paths, labels, class_indices = [], [], {}
|
|
if not os.path.exists(directory):
|
|
return [], [], {}
|
|
for i, cls in enumerate(sorted([d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))])):
|
|
class_indices[cls] = i
|
|
cls_dir = os.path.join(directory, cls)
|
|
for root, _, files in os.walk(cls_dir):
|
|
for file in files:
|
|
if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')):
|
|
file_paths.append(os.path.join(root, file))
|
|
labels.append(i)
|
|
return file_paths, labels, class_indices
|
|
|
|
####################################################################################################
|
|
|
|
def parse_fn(filename, label, image_size, channels, num_classes):
|
|
img = tf.io.read_file(filename)
|
|
img = tf.image.decode_image(img, channels=channels, expand_animations=False)
|
|
img = tf.cond(
|
|
tf.reduce_all(tf.equal(tf.shape(img)[:2], image_size)),
|
|
lambda: tf.cast(img, tf.float32),
|
|
lambda: tf.image.resize(img, image_size))
|
|
img = tf.cast(img, tf.float32) / 255.0
|
|
label = tf.one_hot(label, depth=num_classes)
|
|
return img, label
|
|
|
|
####################################################################################################
|
|
|
|
def load_dataset(data_dir, image_size=(224, 224), val_split=0.01, batch_size=64, color_mode="grayscale"):
|
|
channels = 1 if color_mode == "grayscale" else 3
|
|
file_paths, labels, class_indices = get_file_list(data_dir)
|
|
|
|
if not file_paths:
|
|
raise ValueError(f"No images found in {data_dir}")
|
|
|
|
combined = list(zip(file_paths, labels))
|
|
random.shuffle(combined)
|
|
file_paths[:], labels[:] = zip(*combined)
|
|
|
|
num_val = int(val_split * len(file_paths))
|
|
train_files = file_paths[num_val:]
|
|
train_labels = labels[num_val:]
|
|
val_files = file_paths[:num_val]
|
|
val_labels = labels[:num_val]
|
|
num_classes = len(class_indices)
|
|
|
|
train_ds = tf.data.Dataset.from_tensor_slices((train_files, train_labels))
|
|
train_ds = train_ds.map(lambda f, l: parse_fn(f, l, image_size, channels, num_classes), num_parallel_calls=tf.data.AUTOTUNE)
|
|
train_ds = train_ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
|
|
|
|
val_ds = tf.data.Dataset.from_tensor_slices((val_files, val_labels))
|
|
val_ds = val_ds.map(lambda f, l: parse_fn(f, l, image_size, channels, num_classes), num_parallel_calls=tf.data.AUTOTUNE)
|
|
val_ds = val_ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
|
|
|
|
train_ds.class_indices = class_indices
|
|
return train_ds, val_ds
|
|
|
|
####################################################################################################
|
|
|
|
class CategoricalFocalLoss(losses.Loss):
|
|
def __init__(self, alpha, gamma=2.0, smoothing=0.1, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.alpha = tf.constant(alpha, dtype=tf.float32)
|
|
self.gamma = gamma
|
|
self.smoothing = smoothing
|
|
|
|
def call(self, y_true, y_pred):
|
|
num_classes = tf.cast(tf.shape(y_true)[-1], tf.float32)
|
|
y_true = tf.cast(y_true, tf.float32)
|
|
y_pred = tf.cast(y_pred, tf.float32)
|
|
y_true_smooth = y_true * (1.0 - self.smoothing) + (self.smoothing / (num_classes - 1)) * (1.0 - y_true)
|
|
|
|
log_pt = tf.nn.log_softmax(y_pred)
|
|
pt = tf.exp(log_pt)
|
|
|
|
class_idx = tf.argmax(y_true, axis=-1)
|
|
alpha_factor = tf.gather(self.alpha, class_idx)
|
|
alpha_factor = tf.expand_dims(alpha_factor, axis=-1)
|
|
|
|
focal_weight = alpha_factor * tf.pow(1.0 - pt, self.gamma)
|
|
loss = tf.reduce_sum(-y_true_smooth * log_pt * focal_weight, axis=-1)
|
|
return tf.reduce_mean(loss)
|
|
|
|
def hard_swish(x):
|
|
return x * tf.nn.relu6(x + 3.0) / 6.0
|
|
|
|
def coordinate_attention(x, reduction=4):
|
|
in_channels = x.shape[-1]
|
|
mip = max(8, in_channels // reduction)
|
|
h, w = x.shape[1], x.shape[2]
|
|
|
|
x_h = tf.reduce_mean(x, axis=2, keepdims=True)
|
|
x_w = tf.reduce_mean(x, axis=1, keepdims=True)
|
|
x_w_t = tf.transpose(x_w, perm=[0, 2, 1, 3])
|
|
|
|
y = tf.concat([x_h, x_w_t], axis=1)
|
|
y = layers.Conv2D(mip, 1, use_bias=True, kernel_initializer='he_normal')(y)
|
|
y = layers.BatchNormalization()(y)
|
|
y = layers.Activation(hard_swish)(y)
|
|
|
|
x_h_prime, x_w_prime = tf.split(y, [h, w], axis=1)
|
|
x_w_prime = tf.transpose(x_w_prime, perm=[0, 2, 1, 3])
|
|
|
|
a_h = layers.Conv2D(in_channels, 1, activation='sigmoid', use_bias=True, kernel_initializer='he_normal')(x_h_prime)
|
|
a_w = layers.Conv2D(in_channels, 1, activation='sigmoid', use_bias=True, kernel_initializer='he_normal')(x_w_prime)
|
|
return x * a_h * a_w
|
|
|
|
def mix_kernel_depthwise(x, stride):
|
|
channels = x.shape[-1]
|
|
groups = 3
|
|
c_per_group = channels // groups
|
|
c_last = channels - c_per_group * (groups - 1)
|
|
splits = [c_per_group] * (groups - 1) + [c_last]
|
|
x_splits = tf.split(x, splits, axis=-1)
|
|
|
|
x0 = layers.DepthwiseConv2D(3, strides=stride, padding='same', use_bias=False, depthwise_initializer='he_normal')(x_splits[0])
|
|
x1 = layers.DepthwiseConv2D((1, 7), strides=stride, padding='same', use_bias=False, depthwise_initializer='he_normal')(x_splits[1])
|
|
x2 = layers.DepthwiseConv2D((7, 1), strides=stride, padding='same', use_bias=False, depthwise_initializer='he_normal')(x_splits[2])
|
|
return layers.Concatenate(axis=-1)([x0, x1, x2])
|
|
|
|
def se_block(x, reduction=4):
|
|
in_channels = x.shape[-1]
|
|
y = layers.GlobalAveragePooling2D()(x)
|
|
y = layers.Reshape((1, 1, in_channels))(y)
|
|
y = layers.Conv2D(in_channels // reduction, 1, use_bias=True, kernel_initializer='he_normal')(y)
|
|
y = layers.Activation(tf.nn.gelu)(y)
|
|
y = layers.BatchNormalization(epsilon=1e-3, momentum=0.99)(y)
|
|
y = layers.Conv2D(in_channels, 1, use_bias=True, kernel_initializer='he_normal')(y)
|
|
y = layers.Activation('sigmoid')(y)
|
|
return x * y
|
|
|
|
def mb_conv_block(x, out_channels, stride, expand_ratio, use_se=True):
|
|
in_channels = x.shape[-1]
|
|
use_shortcut = (stride == 1 and in_channels == out_channels)
|
|
hidden_dim = int(round(in_channels * expand_ratio))
|
|
|
|
identity = x
|
|
if expand_ratio != 1:
|
|
x = layers.Conv2D(hidden_dim, 1, use_bias=False, kernel_initializer='he_normal')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.Activation(tf.nn.gelu)(x)
|
|
|
|
x = mix_kernel_depthwise(x, stride)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.Activation(tf.nn.gelu)(x)
|
|
|
|
if use_se:
|
|
x = coordinate_attention(x, reduction=4)
|
|
|
|
x = layers.Conv2D(out_channels, 1, use_bias=False, kernel_initializer='he_normal')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
if use_shortcut:
|
|
return layers.Add()([identity, x])
|
|
return x
|
|
|
|
def build_model(input_shape, num_classes):
|
|
i = layers.Input(shape=input_shape)
|
|
x = layers.Conv2D(32, 3, strides=2, padding='same', use_bias=False, kernel_initializer='he_normal')(i)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.Activation(tf.nn.gelu)(x)
|
|
|
|
config = [
|
|
[32, 48, 1, 3],
|
|
[48, 80, 2, 3],
|
|
[80, 112, 2, 3],
|
|
[112, 160, 2, 3],
|
|
[160, 192, 2, 3],
|
|
[192, 256, 2, 3]]
|
|
|
|
for in_c, out_c, s, ex in config:
|
|
x = mb_conv_block(x, out_c, s, ex)
|
|
|
|
x = layers.Conv2D(512, 1, use_bias=False, kernel_initializer='he_normal')(x)
|
|
x = layers.BatchNormalization()(x)
|
|
x = layers.Activation(tf.nn.gelu)(x)
|
|
x = se_block(x, reduction=4)
|
|
|
|
v_gap = layers.GlobalAveragePooling2D()(x)
|
|
v_gmp = layers.GlobalMaxPooling2D()(x)
|
|
feat = layers.Concatenate(axis=-1)([v_gap, v_gmp])
|
|
|
|
x = layers.Dense(256, kernel_initializer=initializers.RandomNormal(mean=0.0, stddev=0.01))(feat)
|
|
x = layers.Activation(tf.nn.gelu)(x)
|
|
x = layers.Dropout(0.1)(x)
|
|
o = layers.Dense(num_classes, dtype='float32', kernel_initializer=initializers.RandomNormal(mean=0.0, stddev=0.01))(x)
|
|
return models.Model(i, o)
|
|
|
|
####################################################################################################
|
|
|
|
def lr_scheduler(epoch, lr):
|
|
total_epochs, lr_start, lr_end = 25, 1e-3, 1e-6
|
|
return lr_end + (lr_start - lr_end) * (0.5 * (1 + math.cos(epoch / total_epochs * math.pi)) ** 2)
|
|
|
|
class CustomModelCheckpoint(callbacks.Callback):
|
|
def on_epoch_end(self, epoch, logs=False):
|
|
save_name = f"PYAS_Model_Epoch_{epoch + 1}"
|
|
self.model.save(f"{save_name}.h5")
|
|
try:
|
|
spec = (tf.TensorSpec((None,) + self.model.input_shape[1:], tf.float32, name="input"),)
|
|
model_proto, _ = tf2onnx.convert.from_keras(self.model, input_signature=spec, opset=17)
|
|
with open(f"{save_name}.onnx", "wb") as f:
|
|
f.write(model_proto.SerializeToString())
|
|
except Exception as e:
|
|
print(f"ONNX export failed: {e}")
|
|
|
|
####################################################################################################
|
|
|
|
if __name__ == '__main__':
|
|
data_dir = r'.\Image_File_Pefile'
|
|
file_paths, labels, class_indices = get_file_list(data_dir)
|
|
num_classes = len(class_indices)
|
|
|
|
if num_classes == 0:
|
|
print("No classes found. Please check data directory.")
|
|
exit()
|
|
|
|
sample_counts = {}
|
|
for lbl in labels:
|
|
sample_counts[lbl] = sample_counts.get(lbl, 0) + 1
|
|
total_samples = sum(sample_counts.values())
|
|
class_weights_dict = {k: total_samples / (num_classes * v) for k, v in sample_counts.items()}
|
|
alpha_list = [class_weights_dict[i] for i in range(num_classes)]
|
|
|
|
train_ds, val_ds = load_dataset(data_dir)
|
|
input_shape = (224, 224, 1)
|
|
for imgs, _ in train_ds.take(1):
|
|
input_shape = imgs.shape[1:]
|
|
break
|
|
|
|
strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
|
|
with strategy.scope():
|
|
try:
|
|
model = models.load_model('PYAS_Model_Epoch_0.h5', custom_objects={'CategoricalFocalLoss': CategoricalFocalLoss, 'hard_swish': hard_swish})
|
|
print("Load model from disk")
|
|
except:
|
|
print("Creating a new model")
|
|
model = build_model(input_shape, num_classes)
|
|
|
|
print(f"Total parameters: {model.count_params()}")
|
|
optimizer = optimizers.Adam(learning_rate=1e-3)
|
|
criterion = CategoricalFocalLoss(alpha=alpha_list, gamma=2.0)
|
|
model.compile(optimizer=optimizer, loss=criterion, metrics=['accuracy'])
|
|
|
|
model.fit(train_ds, epochs=25, callbacks=[CustomModelCheckpoint(), callbacks.LearningRateScheduler(lr_scheduler)], validation_data=val_ds)
|
|
input('Training Complete')
|