Lab3.py - master - Durbolay/Lab3

Lab3.py
113 строк4 KB
Durbolay
update Lab3.py
19 ноя 2025, 15:51
19 ноя 2025, 15:51b8edb55
# ЛР3: Семантическая сегментация изображений (U-Net, исправленная под Colab)
import tensorflow as tf
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt

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

def normalize_img_mask(img, mask):
    img = tf.cast(img, tf.float32) / 255.0
    mask = tf.cast(mask, tf.int32) - 1  # классы 0,1,2
    return img, mask

def load(datapoint):
    img = tf.image.resize(datapoint['image'], (128, 128))
    mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128), method='nearest')
    return normalize_img_mask(img, mask)

train = ds['train'].map(load, num_parallel_calls=tf.data.AUTOTUNE)
test = ds['test'].map(load, num_parallel_calls=tf.data.AUTOTUNE)

BATCH_SIZE = 16
train_batches = train.cache().shuffle(500).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
test_batches = test.batch(BATCH_SIZE)

# === 2. Определение архитектуры U-Net ===
def conv_block(x, filters):
    x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
    x = tf.keras.layers.Conv2D(filters, 3, padding='same', activation='relu')(x)
    return x

def encoder_block(x, filters):
    c = conv_block(x, filters)
    p = tf.keras.layers.MaxPooling2D((2,2))(c)
    return c, p

def decoder_block(x, skip, filters):
    x = tf.keras.layers.Conv2DTranspose(filters, 2, strides=2, padding='same')(x)
    x = tf.keras.layers.Concatenate()([x, skip])
    x = conv_block(x, filters)
    return x

def build_unet(input_shape=(128,128,3), num_classes=3):
    inputs = tf.keras.Input(input_shape)

    # Encoder
    c1, p1 = encoder_block(inputs, 64)
    c2, p2 = encoder_block(p1, 128)
    c3, p3 = encoder_block(p2, 256)
    c4, p4 = encoder_block(p3, 512)

    # Bridge
    b = conv_block(p4, 1024)

    # Decoder
    d1 = decoder_block(b, c4, 512)
    d2 = decoder_block(d1, c3, 256)
    d3 = decoder_block(d2, c2, 128)
    d4 = decoder_block(d3, c1, 64)

    outputs = tf.keras.layers.Conv2D(num_classes, 1, padding='same', activation='softmax')(d4)
    return tf.keras.Model(inputs, outputs, name="U-Net")

model = build_unet()
model.summary()

# === 3. Определение метрики IoU ===
def mean_iou(y_true, y_pred):
    y_pred = tf.argmax(y_pred, axis=-1)
    y_pred = tf.expand_dims(y_pred, -1)
    y_pred = tf.cast(y_pred, tf.int32)
    y_true = tf.cast(y_true, tf.int32)
    intersect = tf.reduce_sum(tf.cast(y_true == y_pred, tf.float32))
    total = tf.cast(tf.size(y_true), tf.float32)
    return intersect / total

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy', mean_iou])

# === 4. Обучение модели ===
EPOCHS = 5  # можно увеличить при необходимости
history = model.fit(train_batches,
                    validation_data=test_batches,
                    epochs=EPOCHS)

# === 5. Визуализация графиков ===
plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(history.history['loss'], label='train')
plt.plot(history.history['val_loss'], label='val')
plt.title('Loss')
plt.legend()
plt.subplot(1,2,2)
plt.plot(history.history['accuracy'], label='train acc')
plt.plot(history.history['val_accuracy'], label='val acc')
plt.title('Accuracy')
plt.legend()
plt.show()

# === 6. Визуализация предсказаний ===
for images, masks in test_batches.take(1):
    preds = model.predict(images)
    preds = tf.argmax(preds, axis=-1)

    plt.figure(figsize=(12,6))
    for i in range(3):
        plt.subplot(3,3,i*3+1); plt.imshow(images[i]); plt.title('Image'); plt.axis('off')
        plt.subplot(3,3,i*3+2); plt.imshow(tf.squeeze(masks[i]), cmap='jet'); plt.title('True mask'); plt.axis('off')
        plt.subplot(3,3,i*3+3); plt.imshow(preds[i], cmap='jet'); plt.title('Pred mask'); plt.axis('off')
    plt.tight_layout()
    plt.show()
    break