import os os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices=false' os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' import glob libdevice_paths = glob.glob( '/root/miniconda3/envs/deeplab/lib/python3.11/site-packages' '/nvidia/cuda_nvcc/nvvm/libdevice/libdevice.10.bc' ) if libdevice_paths: cuda_dir = os.path.dirname(os.path.dirname(os.path.dirname(libdevice_paths[0]))) os.environ['XLA_FLAGS'] = f'--xla_gpu_cuda_data_dir={cuda_dir}' print(f"[OK] libdevice found: {libdevice_paths[0]}") else: # Fallback: tắt hoàn toàn XLA os.environ['TF_XLA_FLAGS'] = '--tf_xla_enable_xla_devices=false' print("[WARN] libdevice not found, XLA disabled") import warnings warnings.filterwarnings("ignore", category=UserWarning, message=".*input_shape.*") import tensorflow as tf tf.random.set_seed(42) tf.config.optimizer.set_jit(False) tf.config.experimental.set_synchronous_execution(True) from tensorflow.keras import mixed_precision from tensorflow.keras.applications import MobileNetV2 gpus = tf.config.list_physical_devices('GPU') if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) print("[OK] Đã kích hoạt Memory Growth") except RuntimeError as e: print(f"Lỗi khởi tạo GPU: {e}") mixed_precision.set_global_policy('float32') import numpy as np import cv2 import matplotlib.pyplot as plt from sklearn.model_selection import train_test_split from CLMR import CLMRCallback BATCH_SIZE = 8 EPOCHS_STAGE_1 = 20 EPOCHS_STAGE_2 = 80 TARGET_SIZE = (128, 128) NUM_CLASSES = 2 # --- 3. CÁC HÀM TIỀN XỬ LÝ DỮ LIỆU --- def parse_tfrecord(serialized): feature_desc = { 'image': tf.io.FixedLenFeature([], tf.string), 'mask': tf.io.FixedLenFeature([], tf.string), } parsed = tf.io.parse_single_example(serialized, feature_desc) image = tf.image.decode_jpeg(parsed['image'], channels=3) image = tf.cast(image, tf.float32) / 255.0 image = tf.image.resize(image, TARGET_SIZE) image.set_shape([TARGET_SIZE[0], TARGET_SIZE[1], 3]) mask = tf.image.decode_png(parsed['mask'], channels=1) mask = tf.cast(mask, tf.float32) mask = tf.image.resize(mask, TARGET_SIZE, method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) mask.set_shape([TARGET_SIZE[0], TARGET_SIZE[1], 1]) return image, mask def augment(image, mask): """50% crop ngẫu nhiên, 50% crop tập trung vào vùng có class hiếm""" def random_crop(image, mask): combined = tf.concat([image, mask], axis=-1) combined = tf.image.random_flip_left_right(combined) min_crop = tf.cast(TARGET_SIZE[0] * 0.5, tf.int32) crop_size = tf.random.uniform([], minval=min_crop, maxval=TARGET_SIZE[0], dtype=tf.int32) combined = tf.image.random_crop(combined, [crop_size, crop_size, 4]) combined = tf.image.resize(combined, TARGET_SIZE) return combined[..., :3], combined[..., 3:] def focused_crop(image, mask): """Crop vùng có chứa class 2 hoặc 3""" combined = tf.concat([image, mask], axis=-1) combined = tf.image.random_flip_left_right(combined) # Tìm vùng có class hiếm rare_mask = tf.logical_or( tf.equal(mask[..., 0], 2), tf.equal(mask[..., 0], 3) ) rare_indices = tf.where(rare_mask) # Nếu không có class hiếm thì crop ngẫu nhiên has_rare = tf.greater(tf.shape(rare_indices)[0], 0) def crop_around_rare(): # Lấy một điểm ngẫu nhiên trong vùng hiếm idx = tf.random.uniform( [], 0, tf.shape(rare_indices)[0], dtype=tf.int32 ) center = tf.cast(rare_indices[idx], tf.int32) cy, cx = center[0], center[1] h, w = TARGET_SIZE crop_size = tf.cast(h * 0.5, tf.int32) # Tính crop box y1 = tf.clip_by_value(cy - crop_size//2, 0, h - crop_size) x1 = tf.clip_by_value(cx - crop_size//2, 0, w - crop_size) cropped = combined[y1:y1+crop_size, x1:x1+crop_size, :] cropped = tf.image.resize(cropped, TARGET_SIZE) return cropped combined = tf.cond( has_rare, crop_around_rare, lambda: tf.image.resize(combined, TARGET_SIZE) ) return combined[..., :3], combined[..., 3:] # 50% random, 50% focused use_focused = tf.random.uniform([]) > 0.5 image, mask = tf.cond( use_focused, lambda: focused_crop(image, mask), lambda: random_crop(image, mask) ) # Color augmentation image = tf.image.random_brightness(image, max_delta=0.15) image = tf.image.random_contrast(image, lower=0.85, upper=1.15) image = tf.image.random_saturation(image, lower=0.85, upper=1.15) image = tf.clip_by_value(image, 0.0, 1.0) return image, mask def load_tfrecord_dataset(tfrecord_pattern, batch_size=4, training=False): files = tf.data.Dataset.list_files(tfrecord_pattern, shuffle=training) dataset = files.interleave( lambda f: tf.data.TFRecordDataset(f, compression_type='GZIP'), cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE ) dataset = dataset.map(parse_tfrecord, num_parallel_calls=tf.data.AUTOTUNE) if training: dataset = dataset.shuffle(buffer_size=500) dataset = dataset.map(augment, num_parallel_calls=tf.data.AUTOTUNE) dataset = dataset.batch(batch_size, drop_remainder=True) dataset = dataset.prefetch(tf.data.AUTOTUNE) return dataset # --- 5. ĐỊNH NGHĨA MODEL --- def aspp_block(x, filters=128): b1 = tf.keras.layers.Conv2D(filters, 1, padding='same', use_bias=False)(x) b1 = tf.keras.layers.BatchNormalization()(b1) b1 = tf.keras.layers.Activation('relu')(b1) b2 = tf.keras.layers.Conv2D(filters, 3, padding='same', dilation_rate=6, use_bias=False)(x) b2 = tf.keras.layers.BatchNormalization()(b2) b2 = tf.keras.layers.Activation('relu')(b2) out = tf.keras.layers.Concatenate()([b1, b2]) out = tf.keras.layers.Conv2D(filters, 1, use_bias=False)(out) out = tf.keras.layers.BatchNormalization()(out) out = tf.keras.layers.Activation('relu')(out) return out def decodeBlock(prev_layer_input, skip_layer_input, n_filters=32): up = tf.keras.layers.Conv2DTranspose( n_filters, (3,3), # Kernel size strides=(2,2), padding='same', kernel_regularizer=tf.keras.regularizers.l2(1e-4) )(prev_layer_input) merge = tf.keras.layers.concatenate([up, skip_layer_input], axis=3) conv = tf.keras.layers.Conv2D(n_filters, 3, # Kernel size activation='relu', padding='same', kernel_initializer='he_normal', kernel_regularizer=tf.keras.regularizers.l2(1e-4) )(merge) conv = tf.keras.layers.BatchNormalization()(conv) conv = tf.keras.layers.Conv2D(n_filters, 3, # Kernel size activation='relu', padding='same', kernel_initializer='he_normal', kernel_regularizer=tf.keras.regularizers.l2(1e-4) )(conv) conv = tf.keras.layers.BatchNormalization()(conv) return conv from tensorflow.keras.applications import EfficientNetB1 def efficientnet_unet(input_size=(512, 512, 3), n_filter=16, num_classes=5): inputs = tf.keras.Input(shape=input_size) x = tf.keras.layers.Rescaling(scale=255.0)(inputs) backbone = EfficientNetB1( input_tensor=x, include_top=False, weights='imagenet' ) backbone.trainable = False s1 = backbone.get_layer('block2a_expand_activation').output # 256x256 s2 = backbone.get_layer('block3a_expand_activation').output # 128x128 s3 = backbone.get_layer('block4a_expand_activation').output # 64x64 s4 = backbone.get_layer('block6a_expand_activation').output # 32x32 s5 = backbone.get_layer('top_activation').output # 16x16 # Bridge ASPP bridge = aspp_block(s5,filters=n_filter * 8) # Decoder d1 = decodeBlock(bridge,s4,n_filter * 16) # 16→32 d2 = decodeBlock(d1,s3,n_filter * 8) # 32→64 d3 = decodeBlock(d2,s2,n_filter * 4) # 64→128 d4 = decodeBlock(d3,s1,n_filter * 2) # 128→256 x = tf.keras.layers.Conv2DTranspose(n_filter, 3, strides=2, padding='same')(d4) x = tf.keras.layers.BatchNormalization()(x) x = tf.keras.layers.Activation('relu')(x) outputs = tf.keras.layers.Conv2D( num_classes, 1, activation='softmax', dtype='float32' )(x) model = tf.keras.Model(inputs, outputs) model.backbone = backbone return model def deeplab_model(input_size=(512, 512, 3), n_filter=16, num_classes=5): inputs = tf.keras.Input(shape=input_size) x = tf.keras.layers.Rescaling(scale=2.0, offset=-1.0)(inputs) backbone = MobileNetV2( input_tensor=x, include_top=False, weights='imagenet' ) backbone.trainable = False low_level = backbone.get_layer('block_2_add').output # 128x128 high_level = backbone.get_layer('block_13_expand_relu').output # 32x32 aspp_b1 = tf.keras.layers.Conv2D(n_filter, (1,1), padding='same')(high_level) aspp_b1 = tf.keras.layers.BatchNormalization()(aspp_b1) aspp_b1 = tf.keras.layers.Activation('relu')(aspp_b1) aspp_b2 = tf.keras.layers.Conv2D(n_filter, (3,3), dilation_rate=6, padding='same')(high_level) aspp_b2 = tf.keras.layers.BatchNormalization()(aspp_b2) aspp_b2 = tf.keras.layers.Activation('relu')(aspp_b2) aspp_b3 = tf.keras.layers.Conv2D(n_filter, (3,3), dilation_rate=12, padding='same')(high_level) aspp_b3 = tf.keras.layers.BatchNormalization()(aspp_b3) aspp_b3 = tf.keras.layers.Activation('relu')(aspp_b3) aspp_b4 = tf.keras.layers.Conv2D(n_filter, (3,3), dilation_rate=18, padding='same')(high_level) aspp_b4 = tf.keras.layers.BatchNormalization()(aspp_b4) aspp_b4 = tf.keras.layers.Activation('relu')(aspp_b4) aspp_b5 = tf.keras.layers.GlobalAveragePooling2D(keepdims=True)(high_level) aspp_b5 = tf.keras.layers.Conv2D(n_filter, (1,1), padding='same')(aspp_b5) aspp_b5 = tf.keras.layers.BatchNormalization()(aspp_b5) aspp_b5 = tf.keras.layers.Activation('relu')(aspp_b5) #aspp_b5 = tf.keras.layers.UpSampling2D(size=(32, 32), interpolation='bilinear')(aspp_b5) pool_size = (high_level.shape[1], high_level.shape[2]) aspp_b5 = tf.keras.layers.Lambda(lambda x: tf.image.resize(x, pool_size, method='bilinear'))(aspp_b5) aspp_out = tf.keras.layers.Concatenate()([aspp_b1, aspp_b2, aspp_b3, aspp_b4, aspp_b5]) aspp_out = tf.keras.layers.Conv2D(n_filter, (1,1), padding='same')(aspp_out) aspp_out = tf.keras.layers.BatchNormalization()(aspp_out) aspp_out = tf.keras.layers.Activation('relu')(aspp_out) high_branch = tf.keras.layers.UpSampling2D(size=(4,4), interpolation='bilinear')(aspp_out) # 32→128 low_branch = tf.keras.layers.Conv2D(48, 1, use_bias=False)(low_level) low_branch = tf.keras.layers.BatchNormalization()(low_branch) low_branch = tf.keras.layers.Activation('relu')(low_branch) encode = tf.keras.layers.Concatenate()([low_branch, high_branch]) # 128x128 encode = tf.keras.layers.Conv2D(n_filter, (3,3), padding='same')(encode) encode = tf.keras.layers.BatchNormalization()(encode) encode = tf.keras.layers.Activation('relu')(encode) encode = tf.keras.layers.Conv2D(n_filter, (3,3), padding='same')(encode) encode = tf.keras.layers.BatchNormalization()(encode) encode = tf.keras.layers.Activation('relu')(encode) encode = tf.keras.layers.UpSampling2D(size=(4,4), interpolation='bilinear')(encode) # 128→512 outputs = tf.keras.layers.Conv2D(num_classes, 1, activation='softmax', dtype='float32')(encode) model = tf.keras.Model(inputs, outputs) model.backbone = backbone return model # --- 6. LOSS & METRICS & OPTIMIZER & CALLBACKS--- class SparseMeanIoU(tf.keras.metrics.MeanIoU): def update_state(self, y_true, y_pred, sample_weight=None): y_pred = tf.argmax(y_pred, axis=-1) if len(y_true.shape) == 4: y_true = tf.squeeze(y_true, axis=-1) y_true = tf.cast(y_true, tf.int64) y_pred = tf.cast(y_pred, tf.int64) y_true = tf.where(y_true == 5, tf.cast(1, tf.int64), y_true) y_true = tf.clip_by_value(y_true, 0, self.num_classes - 1) return super().update_state(y_true, y_pred, sample_weight) def tversky_loss(y_true, y_pred, alpha=0.7, beta=0.3): y_true_sparse = tf.cast(y_true[..., 0], tf.int32) y_true_oh = tf.one_hot(y_true_sparse, depth=5) tp = tf.reduce_sum(y_true_oh * y_pred, axis=[1, 2]) fp = tf.reduce_sum((1 - y_true_oh) * y_pred, axis=[1, 2]) fn = tf.reduce_sum(y_true_oh * (1 - y_pred), axis=[1, 2]) tversky = (tp + 1e-6) / (tp + alpha * fn + beta * fp + 1e-6) return 1.0 - tf.reduce_mean(tversky) lr_schedule = tf.keras.optimizers.schedules.CosineDecayRestarts( initial_learning_rate=1e-4, # LR ban đầu first_decay_steps=1000, # Số steps mỗi chu kỳ restart t_mul=2.0, # Mỗi chu kỳ sau dài gấp đôi m_mul=0.9, # LR đỉnh giảm 10% sau mỗi restart alpha=1e-6, # LR tối thiểu (không về 0) ) # --- 7. TRAINING --- os.makedirs('./model', exist_ok=True) # Mô Hình deeplabv3 print("Đang khởi tạo Dataset trên CPU...") train_data = load_tfrecord_dataset('./tfrecord/train_*.tfrecord', batch_size=4, training=True) val_data = load_tfrecord_dataset('./tfrecord/val_*.tfrecord', batch_size=4, training=False) test_data = load_tfrecord_dataset('./tfrecord/test_*.tfrecord', batch_size=4,training=False) print("[OK] Dataset đã sẵn sàng") model = deeplab_model(input_size=(128, 128, 3), n_filter=32) print("[OK] Đã dựng xong mô hình") model.compile( optimizer=tf.keras.optimizers.Adam( learning_rate=lr_schedule, beta_1=0.9, # Momentum (default, giữ nguyên) beta_2=0.999, # RMSprop (default, giữ nguyên) epsilon=1e-7, # Tránh chia 0 (default) clipnorm=1.0, # Gradient clipping: tránh gradient explosion ), loss=tversky_loss, metrics=['accuracy', SparseMeanIoU(num_classes=2, name='mean_iou')] ) # Giai đoạn 1 print("Starting Stage 1...") history_1 = model.fit(train_data, epochs=EPOCHS_STAGE_1, validation_data=val_data) plt.figure(figsize=(12,6)) plt.subplot(1,3,1) plt.plot(history_1.history['accuracy'], label='accuracy') plt.plot(history_1.history['val_accuracy'], label='val_accuracy') plt.title('Model accuracy'); plt.xlabel('Epoch'); plt.ylabel('accuracy'); plt.legend() plt.subplot(1,3,2) plt.plot(history_1.history['mean_iou'], label='mean_iou') plt.plot(history_1.history['val_mean_iou'], label='val_mean_iou') plt.title('Model mean iou'); plt.xlabel('Epoch'); plt.ylabel('mean iou'); plt.legend() plt.subplot(1,3,3) plt.plot(history_1.history['loss'], label='loss') plt.plot(history_1.history['val_loss'], label='val_loss') plt.title('Model loss'); plt.xlabel('Epoch'); plt.ylabel('loss'); plt.legend() plt.savefig('history_stage1.png'); plt.close() # Giai đoạn 2 model.backbone.trainable = True model.compile( optimizer=tf.keras.optimizers.Adam( learning_rate=lr_schedule, beta_1=0.9, # Momentum (default, giữ nguyên) beta_2=0.999, # RMSprop (default, giữ nguyên) epsilon=1e-7, # Tránh chia 0 (default) clipnorm=1.0, # Gradient clipping: tránh gradient explosion ), loss=tversky_loss, metrics=['accuracy', SparseMeanIoU(num_classes=2, name='mean_iou')] ) print("Starting Stage 2 (Fine-tuning)...") checkpoint = tf.keras.callbacks.ModelCheckpoint( './model/best_deeplab_eff.keras', save_best_only=True, monitor='val_mean_iou', mode='max' ) history_2 = model.fit(train_data, epochs=EPOCHS_STAGE_2, validation_data=val_data, callbacks=[checkpoint]) plt.figure(figsize=(12,6)) plt.subplot(1,3,1) plt.plot(history_2.history['accuracy'], label='accuracy') plt.plot(history_2.history['val_accuracy'], label='val_accuracy') plt.title('Model accuracy'); plt.xlabel('Epoch'); plt.ylabel('accuracy'); plt.legend() plt.subplot(1,3,2) plt.plot(history_2.history['mean_iou'], label='mean_iou') plt.plot(history_2.history['val_mean_iou'], label='val_mean_iou') plt.title('Model mean iou'); plt.xlabel('Epoch'); plt.ylabel('mean iou'); plt.legend() plt.subplot(1,3,3) plt.plot(history_2.history['loss'], label='loss') plt.plot(history_2.history['val_loss'], label='val_loss') plt.title('Model loss'); plt.xlabel('Epoch'); plt.ylabel('loss'); plt.legend() plt.savefig('./model/history_stage2_deeplab_mobilenet.png'); plt.close() model.save('./model/deeplab_mobilenet.h5') print("Training Complete. Model Deeplab saved.") #Mô Hình effnet model_eff = efficientnet_unet(input_size=(128, 128, 3), n_filter=32) print("[OK] Đã dựng xong mô hình effnet") model_eff.compile( optimizer=tf.keras.optimizers.Adam( learning_rate=lr_schedule, beta_1=0.9, # Momentum (default, giữ nguyên) beta_2=0.999, # RMSprop (default, giữ nguyên) epsilon=1e-7, # Tránh chia 0 (default) clipnorm=1.0, # Gradient clipping: tránh gradient explosion ), loss=tversky_loss, metrics=['accuracy', SparseMeanIoU(num_classes=2, name='mean_iou')] ) # Giai đoạn 1 print("Starting Stage 1...") history_1 = model_eff.fit(train_data, epochs=EPOCHS_STAGE_1, validation_data=val_data) plt.figure(figsize=(12,6)) plt.subplot(1,3,1) plt.plot(history_1.history['accuracy'], label='accuracy') plt.plot(history_1.history['val_accuracy'], label='val_accuracy') plt.title('Model accuracy'); plt.xlabel('Epoch'); plt.ylabel('accuracy'); plt.legend() plt.subplot(1,3,2) plt.plot(history_1.history['mean_iou'], label='mean_iou') plt.plot(history_1.history['val_mean_iou'], label='val_mean_iou') plt.title('Model mean iou'); plt.xlabel('Epoch'); plt.ylabel('mean iou'); plt.legend() plt.subplot(1,3,3) plt.plot(history_1.history['loss'], label='loss') plt.plot(history_1.history['val_loss'], label='val_loss') plt.title('Model loss'); plt.xlabel('Epoch'); plt.ylabel('loss'); plt.legend() plt.savefig('history_stage1(effnet).png'); plt.close() # Giai đoạn 2 model_eff.backbone.trainable = True model_eff.compile( optimizer=tf.keras.optimizers.Adam( learning_rate=lr_schedule, beta_1=0.9, # Momentum (default, giữ nguyên) beta_2=0.999, # RMSprop (default, giữ nguyên) epsilon=1e-7, # Tránh chia 0 (default) clipnorm=1.0, # Gradient clipping: tránh gradient explosion ), loss=tversky_loss, metrics=['accuracy', SparseMeanIoU(num_classes=2, name='mean_iou')] ) print("Starting Stage 2 (Fine-tuning)...") checkpoint = tf.keras.callbacks.ModelCheckpoint( './model/best_eff_net.keras', save_best_only=True, monitor='val_mean_iou', mode='max' ) history_2 = model_eff.fit(train_data, epochs=EPOCHS_STAGE_2, validation_data=val_data, callbacks=[checkpoint]) plt.figure(figsize=(12,6)) plt.subplot(1,3,1) plt.plot(history_2.history['accuracy'], label='accuracy') plt.plot(history_2.history['val_accuracy'], label='val_accuracy') plt.title('Model accuracy'); plt.xlabel('Epoch'); plt.ylabel('accuracy'); plt.legend() plt.subplot(1,3,2) plt.plot(history_2.history['mean_iou'], label='mean_iou') plt.plot(history_2.history['val_mean_iou'], label='val_mean_iou') plt.title('Model mean iou'); plt.xlabel('Epoch'); plt.ylabel('mean iou'); plt.legend() plt.subplot(1,3,3) plt.plot(history_2.history['loss'], label='loss') plt.plot(history_2.history['val_loss'], label='val_loss') plt.title('Model loss'); plt.xlabel('Epoch'); plt.ylabel('loss'); plt.legend() plt.savefig('./model/history_stage2_eff_net.png'); plt.close() model_eff.save('./model/eff_net.h5') print("Training Complete. Model Effnet saved.") loss, accuracy, mean_iou = model.evaluate(test_data) print(f"Loss DeepLabV3+ :{loss}") print(f"Accuracy DeepLabV3+ :{accuracy}") print(f"Mean IoU DeepLabV3+ :{mean_iou}") loss, accuracy, mean_iou = model_eff.evaluate(test_data) print(f"Loss Effnet :{loss}") print(f"Accuracy Effnet :{accuracy}") print(f"Mean IoU Effnet :{mean_iou}")