Spaces:
Sleeping
Sleeping
File size: 8,743 Bytes
7a5bb5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | {
"nbformat": 4,
"nbformat_minor": 4,
"metadata": {},
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# DeepOceans - Satellite Semantic Segmentation Training\nMake sure to go to **Runtime > Change runtime type > T4 GPU** before running."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 1. Mount Google Drive\nfrom google.colab import drive\ndrive.mount('/content/drive')\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 2. Imports and Setup\nimport os\nimport glob\nimport tensorflow as tf\nfrom tensorflow.keras import layers, models, backend as K\nfrom tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau\nfrom sklearn.model_selection import train_test_split\n\nprint(\"TensorFlow Version:\", tf.__version__)\nprint(\"GPU Available:\", tf.config.list_physical_devices('GPU'))\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 3. U-Net Architecture Construction\ndef conv_block(input_tensor, num_filters):\n x = layers.Conv2D(num_filters, (3, 3), padding=\"same\")(input_tensor)\n x = layers.BatchNormalization()(x)\n x = layers.Activation(\"relu\")(x)\n x = layers.Conv2D(num_filters, (3, 3), padding=\"same\")(x)\n x = layers.BatchNormalization()(x)\n x = layers.Activation(\"relu\")(x)\n return x\n\ndef encoder_block(input_tensor, num_filters):\n x = conv_block(input_tensor, num_filters)\n p = layers.MaxPooling2D((2, 2))(x)\n return x, p\n\ndef decoder_block(input_tensor, concat_tensor, num_filters):\n x = layers.Conv2DTranspose(num_filters, (2, 2), strides=(2, 2), padding=\"same\")(input_tensor)\n x = layers.concatenate([x, concat_tensor])\n x = conv_block(x, num_filters)\n return x\n\ndef build_unet(input_shape=(256, 256, 3)):\n inputs = layers.Input(shape=input_shape)\n e1, p1 = encoder_block(inputs, 64)\n e2, p2 = encoder_block(p1, 128)\n e3, p3 = encoder_block(p2, 256)\n e4, p4 = encoder_block(p3, 512)\n b = conv_block(p4, 1024)\n d1 = decoder_block(b, e4, 512)\n d2 = decoder_block(d1, e3, 256)\n d3 = decoder_block(d2, e2, 128)\n d4 = decoder_block(d3, e1, 64)\n outputs = layers.Conv2D(1, (1, 1), padding=\"same\", activation=\"sigmoid\")(d4)\n model = models.Model(inputs, outputs, name=\"U-Net\")\n return model\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 4. Custom Loss Functions & Metrics (Dice + BCE)\ndef dice_coef(y_true, y_pred, smooth=1e-6):\n y_true_f = K.flatten(y_true)\n y_pred_f = K.flatten(y_pred)\n intersection = K.sum(y_true_f * y_pred_f)\n return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)\n\ndef dice_loss(y_true, y_pred):\n return 1 - dice_coef(y_true, y_pred)\n\ndef bce_dice_loss(y_true, y_pred):\n bce = tf.keras.losses.binary_crossentropy(y_true, y_pred)\n bce = K.mean(bce) \n return bce + dice_loss(y_true, y_pred)\n\ndef iou_metric(y_true, y_pred, smooth=1e-6):\n y_true_f = K.flatten(y_true)\n y_pred_f = K.flatten(y_pred)\n intersection = K.sum(y_true_f * y_pred_f)\n union = K.sum(y_true_f) + K.sum(y_pred_f) - intersection\n return (intersection + smooth) / (union + smooth)\n\ndef get_metrics():\n return [\n dice_coef,\n iou_metric,\n tf.keras.metrics.Precision(name='precision'),\n tf.keras.metrics.Recall(name='recall')\n ]\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 5. Data Pipeline & Augmentation\nIMG_SIZE = (256, 256)\n\ndef decode_image(image_file):\n image = tf.io.read_file(image_file)\n image = tf.image.decode_image(image, channels=3, expand_animations=False)\n image = tf.image.resize(image, IMG_SIZE)\n image = tf.cast(image, tf.float32) / 255.0\n return image\n\ndef decode_mask(mask_file):\n mask = tf.io.read_file(mask_file)\n mask = tf.image.decode_image(mask, channels=1, expand_animations=False)\n mask = tf.image.resize(mask, IMG_SIZE)\n mask = tf.cast(mask, tf.float32) / 255.0\n mask = tf.math.round(mask) \n return mask\n\ndef process_path(image_path, mask_path):\n image = decode_image(image_path)\n mask = decode_mask(mask_path)\n return image, mask\n\ndef augment(image, mask):\n if tf.random.uniform(()) > 0.5:\n image = tf.image.flip_left_right(image)\n mask = tf.image.flip_left_right(mask)\n if tf.random.uniform(()) > 0.5:\n image = tf.image.flip_up_down(image)\n mask = tf.image.flip_up_down(mask)\n image = tf.image.random_brightness(image, max_delta=0.2)\n image = tf.clip_by_value(image, 0.0, 1.0)\n return image, mask\n\ndef get_dataset(image_paths, mask_paths, batch_size=16, is_train=True):\n dataset = tf.data.Dataset.from_tensor_slices((image_paths, mask_paths))\n if is_train:\n dataset = dataset.shuffle(buffer_size=1000)\n dataset = dataset.map(process_path, num_parallel_calls=tf.data.AUTOTUNE)\n if is_train:\n dataset = dataset.map(augment, num_parallel_calls=tf.data.AUTOTUNE)\n dataset = dataset.batch(batch_size)\n dataset = dataset.prefetch(tf.data.AUTOTUNE)\n return dataset\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 6. Execute Training\n\n# ==============================================================================\n# \u26a0\ufe0f ACTION REQUIRED: Verify this points to the folder containing your dataset\n# (Ensure there is an 'images/' and 'masks/' folder inside this directory)\n# ==============================================================================\nDATA_DIR = '/content/drive/MyDrive/1j35RM-5uZbTOfDv0mBTeRRmOUPvvg28t'\n# ==============================================================================\n\nEPOCHS = 50\nBATCH_SIZE = 16\nLEARNING_RATE = 1e-4\n\ndef get_file_paths(data_dir):\n image_dir = os.path.join(data_dir, 'images')\n mask_dir = os.path.join(data_dir, 'masks')\n \n image_paths = glob.glob(os.path.join(image_dir, '*.*'))\n image_paths.sort()\n \n final_img_paths = []\n final_mask_paths = []\n for img in image_paths:\n base = os.path.splitext(os.path.basename(img))[0]\n masks = glob.glob(os.path.join(mask_dir, f\"{base}.*\"))\n if len(masks) > 0:\n final_img_paths.append(img)\n final_mask_paths.append(masks[0])\n\n return final_img_paths, final_mask_paths\n\nprint(\"Hunting for dataset at folder:\", DATA_DIR)\nimg_paths, mask_paths = get_file_paths(DATA_DIR)\nprint(f\"Found {len(img_paths)} valid image/mask pairs.\\n\")\n\nif len(img_paths) == 0:\n print(\"\u274c ERROR: No images found. Check that DATA_DIR contains your 'images' and 'masks' folders.\")\nelse:\n # 80/20 train/validation split\n train_x, val_x, train_y, val_y = train_test_split(img_paths, mask_paths, test_size=0.2, random_state=42)\n \n train_dataset = get_dataset(train_x, train_y, batch_size=BATCH_SIZE, is_train=True)\n val_dataset = get_dataset(val_x, val_y, batch_size=BATCH_SIZE, is_train=False)\n \n model = build_unet(input_shape=(256, 256, 3))\n optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)\n \n model.compile(optimizer=optimizer, loss=bce_dice_loss, metrics=get_metrics())\n \n # Save directly to the root of your Google Drive so it's easy to download\n model_save_path = '/content/drive/MyDrive/oil_spill_unet_best.keras'\n \n callbacks = [\n ModelCheckpoint(model_save_path, verbose=1, save_best_only=True, monitor='val_loss'),\n EarlyStopping(monitor='val_loss', patience=15, restore_best_weights=True),\n ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=5, verbose=1, min_lr=1e-6)\n ]\n \n print(\"\ud83d\ude80 Initializing Training Pipeline...\")\n history = model.fit(\n train_dataset,\n validation_data=val_dataset,\n epochs=EPOCHS,\n callbacks=callbacks\n )\n \n print(\"\\n\u2705 Training Complete!\")\n print(f\"Your trained model file has been successfully saved to: {model_save_path}\")\n print(\"\u2b07\ufe0f You can now download this file and move it to your local 'model/saved_models/' folder.\")\n"
]
}
]
} |