CodeLab3 - master - soft_3/Lab_CIS

CodeLab3
225 строк8 KB
soft_3
create CodeLab3
24 окт 2025, 12:43
24 окт 2025, 12:4303b627f
import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import classification_report
import os

# Проверка GPU
print("GPU доступен:", tf.config.list_physical_devices('GPU'))

# 1. Загрузка и предобработка данных
dataset, info = tfds.load('oxford_iiit_pet', with_info=True, as_supervised=False)

NUM_CLASSES = 3  # background, foreground (pet), outline
IMG_SIZE = (128, 128)
BATCH_SIZE = 32

def normalize_and_resize(image, mask):
    image = tf.image.resize(image, IMG_SIZE)
    mask = tf.image.resize(mask, IMG_SIZE, method='nearest')
    image = tf.cast(image, tf.float32) / 255.0
    mask = tf.cast(mask, tf.int32) - 1  # значения маски: 1,2,3 → 0,1,2
    return image, mask

def preprocess_train(example):
    image, mask = example['image'], example['segmentation_mask']
    # Аугментация: случайное горизонтальное отражение
    if tf.random.uniform(()) > 0.5:
        image = tf.image.flip_left_right(image)
        mask = tf.image.flip_left_right(mask)
    return normalize_and_resize(image, mask)

def preprocess_test(example):
    image, mask = example['image'], example['segmentation_mask']
    return normalize_and_resize(image, mask)

# Создание датасетов
train_ds = dataset['train'].map(preprocess_train, num_parallel_calls=tf.data.AUTOTUNE)
test_ds = dataset['test'].map(preprocess_test, num_parallel_calls=tf.data.AUTOTUNE)

# Настройка производительности
train_ds = train_ds.cache().shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
test_ds = test_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# Визуализация одного примера
def show_example(ds):
    for image, mask in ds.take(1):
        plt.figure(figsize=(12, 4))
        plt.subplot(1, 3, 1)
        plt.imshow(image[0])
        plt.title("Изображение")
        plt.axis('off')

        plt.subplot(1, 3, 2)
        plt.imshow(mask[0][:, :, 0], cmap='jet', vmin=0, vmax=NUM_CLASSES-1)
        plt.title("Истинная маска")
        plt.axis('off')

        plt.subplot(1, 3, 3)
        plt.imshow(image[0] * 0.6 + plt.get_cmap('jet')(mask[0][:, :, 0]/(NUM_CLASSES-1))[:, :, :3] * 0.4)
        plt.title("Наложение")
        plt.axis('off')
        plt.show()

show_example(train_ds)

# 2. Реализация U-Net с нуля
def double_conv_block(x, n_filters):
    x = tf.keras.layers.Conv2D(n_filters, 3, padding="same", activation="relu")(x)
    x = tf.keras.layers.Conv2D(n_filters, 3, padding="same", activation="relu")(x)
    return x

def downsample_block(x, n_filters):
    f = double_conv_block(x, n_filters)
    p = tf.keras.layers.MaxPool2D(2)(f)
    return f, p

def upsample_block(x, conv_features, n_filters):
    x = tf.keras.layers.Conv2DTranspose(n_filters, 3, 2, padding="same")(x)
    x = tf.keras.layers.concatenate([x, conv_features])
    x = double_conv_block(x, n_filters)
    return x

def build_unet_model(input_shape, num_classes):
    inputs = tf.keras.layers.Input(shape=input_shape)

    # Encoder (downsampling)
    f1, p1 = downsample_block(inputs, 64)
    f2, p2 = downsample_block(p1, 128)
    f3, p3 = downsample_block(p2, 256)
    f4, p4 = downsample_block(p3, 512)

    # Bottleneck
    bottleneck = double_conv_block(p4, 1024)

    # Decoder (upsampling)
    u6 = upsample_block(bottleneck, f4, 512)
    u7 = upsample_block(u6, f3, 256)
    u8 = upsample_block(u7, f2, 128)
    u9 = upsample_block(u8, f1, 64)

    # Output layer
    outputs = tf.keras.layers.Conv2D(num_classes, 1, activation="softmax")(u9)

    model = tf.keras.Model(inputs, outputs)
    return model

model = build_unet_model(input_shape=(*IMG_SIZE, 3), num_classes=NUM_CLASSES)
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

model.summary()

# 3. Обучение модели
EPOCHS = 20

history = model.fit(
    train_ds,
    validation_data=test_ds,
    epochs=EPOCHS,
    verbose=1
)

# 4. Визуализация обучения
def plot_history(history):
    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']
    loss = history.history['loss']
    val_loss = history.history['val_loss']

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(acc, label='Training Accuracy')
    plt.plot(val_acc, label='Validation Accuracy')
    plt.legend()
    plt.title('Accuracy')

    plt.subplot(1, 2, 2)
    plt.plot(loss, label='Training Loss')
    plt.plot(val_loss, label='Validation Loss')
    plt.legend()
    plt.title('Loss')
    plt.show()

plot_history(history)

# 5. Оценка и визуализация результатов
def visualize_predictions(model, dataset, num=3):
    for image_batch, mask_batch in dataset.take(1):
        pred_mask = model.predict(image_batch)
        pred_mask = tf.argmax(pred_mask, axis=-1)

        plt.figure(figsize=(15, 5 * num))
        for i in range(min(num, image_batch.shape[0])):
            plt.subplot(num, 3, i * 3 + 1)
            plt.imshow(image_batch[i])
            plt.title("Изображение")
            plt.axis('off')

            plt.subplot(num, 3, i * 3 + 2)
            plt.imshow(mask_batch[i][:, :, 0], cmap='jet', vmin=0, vmax=NUM_CLASSES-1)
            plt.title("Истинная маска")
            plt.axis('off')

            plt.subplot(num, 3, i * 3 + 3)
            plt.imshow(pred_mask[i], cmap='jet', vmin=0, vmax=NUM_CLASSES-1)
            plt.title("Предсказанная маска")
            plt.axis('off')
        plt.tight_layout()
        plt.show()
        break

visualize_predictions(model, test_ds)

# 6. Метрики IoU и Dice
def mean_iou(y_true, y_pred, num_classes=3):
    y_pred = tf.argmax(y_pred, axis=-1)  # → (batch, H, W)
    y_true = tf.cast(tf.squeeze(y_true, axis=-1), tf.int32)  # → (batch, H, W)
    y_pred = tf.cast(y_pred, tf.int32)   # → (batch, H, W)

    iou = []
    for cls in range(num_classes):
        true_cls = tf.equal(y_true, cls)  # → (batch, H, W)
        pred_cls = tf.equal(y_pred, cls)  # → (batch, H, W)
        intersection = tf.reduce_sum(tf.cast(tf.logical_and(true_cls, pred_cls), tf.float32))
        union = tf.reduce_sum(tf.cast(tf.logical_or(true_cls, pred_cls), tf.float32))
        iou.append(tf.where(union == 0, 1.0, intersection / (union + 1e-7)))
    return tf.reduce_mean(iou)

def dice_coeff(y_true, y_pred, num_classes=3):
    y_pred = tf.argmax(y_pred, axis=-1)  # → (batch, H, W)
    y_true = tf.cast(tf.squeeze(y_true, axis=-1), tf.int32)  # → (batch, H, W)
    y_pred = tf.cast(y_pred, tf.int32)   # → (batch, H, W)

    dice = []
    for cls in range(num_classes):
        true_cls = tf.equal(y_true, cls)
        pred_cls = tf.equal(y_pred, cls)
        intersection = tf.reduce_sum(tf.cast(tf.logical_and(true_cls, pred_cls), tf.float32))
        sum_ = tf.reduce_sum(tf.cast(true_cls, tf.float32)) + tf.reduce_sum(tf.cast(pred_cls, tf.float32))
        dice.append(tf.where(sum_ == 0, 1.0, (2. * intersection) / (sum_ + 1e-7)))
    return tf.reduce_mean(dice)

# Оценка на тестовом наборе
test_loss, test_acc = model.evaluate(test_ds, verbose=0)
print(f"\nТестовая точность: {test_acc:.4f}")

# Ручной расчёт IoU и Dice
all_true, all_pred = [], []
for image_batch, mask_batch in test_ds.take(-1):
    pred_batch = model.predict(image_batch, verbose=0)
    all_true.append(mask_batch)
    all_pred.append(pred_batch)

all_true = tf.concat(all_true, axis=0)
all_pred = tf.concat(all_pred, axis=0)

iou_score = mean_iou(all_true, all_pred)
dice_score = dice_coeff(all_true, all_pred)

print(f"Mean IoU: {iou_score:.4f}")
print(f"Dice Coefficient: {dice_score:.4f}")