diff --git a/training_pipeline/deepwatermap_model.py b/training_pipeline/deepwatermap_model.py new file mode 100644 index 0000000..74e9b2d --- /dev/null +++ b/training_pipeline/deepwatermap_model.py @@ -0,0 +1,281 @@ +""" +4-Channel DeepWaterMap Architecture, Loss Functions, and Dataset Loader for Planet Labs Data + +Implements the DeepWaterMap architecture adapted for 4-band Planet Labs satellite imagery (Blue, Green, Red, NIR). +Includes BCEDiceLoss and comprehensive online spatial & photometric data augmentations. +""" + +import os +import glob +import re +import numpy as np +import rasterio as rio +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset + +class BCEDiceLoss(nn.Module): + """ + Combined Binary Cross Entropy and Dice Loss for coastal boundary segmentation. + Penalizes both pixel-level classification errors and overall mask boundary dissimilarity. + """ + def __init__(self, bce_weight=0.5): + super().__init__() + self.bce = nn.BCEWithLogitsLoss() + self.bce_weight = bce_weight + + def forward(self, logits, targets): + bce_loss = self.bce(logits, targets) + + probs = torch.sigmoid(logits) + smooth = 1.0 + intersection = (probs * targets).sum(dim=(2, 3)) + union = probs.sum(dim=(2, 3)) + targets.sum(dim=(2, 3)) + dice_loss = 1.0 - (2.0 * intersection + smooth) / (union + smooth) + dice_loss = dice_loss.mean() + + return self.bce_weight * bce_loss + (1.0 - self.bce_weight) * dice_loss + +class ConvBlockDWM(nn.Module): + """Convolution block with BatchNorm and optional ReLU activation.""" + def __init__(self, in_c, out_c, k_size, stride=1, use_relu=True): + super().__init__() + padding = k_size // 2 + self.conv = nn.Conv2d(in_c, out_c, kernel_size=k_size, stride=stride, padding=padding, bias=False) + self.bn = nn.BatchNorm2d(out_c, eps=1e-3, momentum=0.01) + self.use_relu = use_relu + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + if self.use_relu: + x = F.relu(x) + return x + +class DownscalingUnitDWM(nn.Module): + """Downscaling unit with residual addition.""" + def __init__(self, in_c, out_c): + super().__init__() + self.c1 = ConvBlockDWM(in_c, out_c, k_size=5, stride=2, use_relu=True) + self.c2 = ConvBlockDWM(out_c, out_c, k_size=3, stride=1, use_relu=True) + + def forward(self, x): + x1 = self.c1(x) + x2 = self.c2(x1) + return x1 + x2 + +class UpscalingUnitDWM(nn.Module): + """Upscaling unit using PixelShuffle (sub-pixel convolution).""" + def __init__(self, in_c, out_c): + super().__init__() + self.pixel_shuffle = nn.PixelShuffle(2) + self.c1 = ConvBlockDWM(in_c // 4, out_c, k_size=3, stride=1, use_relu=True) + self.c2 = ConvBlockDWM(out_c, out_c, k_size=3, stride=1, use_relu=True) + + def forward(self, x): + x = self.pixel_shuffle(x) + x1 = self.c1(x) + x2 = self.c2(x1) + return x1 + x2 + +class BottleneckUnitDWM(nn.Module): + """Bottleneck unit with residual connection.""" + def __init__(self, c): + super().__init__() + self.c1 = ConvBlockDWM(c, c, k_size=3, stride=1, use_relu=True) + self.c2 = ConvBlockDWM(c, c, k_size=3, stride=1, use_relu=True) + + def forward(self, x): + x1 = self.c1(x) + x2 = self.c2(x1) + return x1 + x2 + +class DeepWaterMap4Chan(nn.Module): + """ + 4-Channel PyTorch DeepWaterMap model for Planet Labs satellite data (RGB + NIR). + """ + def __init__(self, in_channels=4): + super().__init__() + self.first_layer = ConvBlockDWM(in_channels, 4, k_size=1, stride=1, use_relu=False) + self.down1 = DownscalingUnitDWM(4, 16) + self.down2 = DownscalingUnitDWM(16, 64) + self.down3 = DownscalingUnitDWM(64, 256) + self.down4 = DownscalingUnitDWM(256, 1024) + + self.bottleneck = BottleneckUnitDWM(1024) + + self.up1 = UpscalingUnitDWM(1024, 256) + self.up2 = UpscalingUnitDWM(256, 64) + self.up3 = UpscalingUnitDWM(64, 16) + self.up4 = UpscalingUnitDWM(16, 4) + + self.last_layer = ConvBlockDWM(4, 1, k_size=1, stride=1, use_relu=False) + + def forward(self, x): + skips = [] + x0 = self.first_layer(x) + skips.append(x0) + + x1 = self.down1(x0) + skips.append(x1) + + x2 = self.down2(x1) + skips.append(x2) + + x3 = self.down3(x2) + skips.append(x3) + + x4 = self.down4(x3) + skips.append(x4) + + b = self.bottleneck(x4) + + d1 = b + skips.pop() + u1 = self.up1(d1) + + d2 = u1 + skips.pop() + u2 = self.up2(d2) + + d3 = u2 + skips.pop() + u3 = self.up3(d3) + + d4 = u3 + skips.pop() + u4 = self.up4(d4) + + d5 = u4 + skips.pop() + out = self.last_layer(d5) + return out + +class SegmentationDataset4Chan(Dataset): + """ + Dataset class for 4-channel (RGB + NIR) Planet Labs satellite imagery and binary NDWI water masks. + Supports comprehensive spatial (flips, rotations, random scale/crop) and photometric (brightness jitter, noise) augmentations. + """ + def __init__(self, data_dir, image_size=(256, 256), is_train=False): + self.data_dir = data_dir + self.image_size = image_size + self.is_train = is_train + self.image_mask_pairs = self._find_pairs() + + def _find_pairs(self): + pairs = [] + image_files = glob.glob(os.path.join(self.data_dir, "*.tif")) + image_files = [f for f in image_files if "_concatenated_ndwi_mask_" not in os.path.basename(f)] + + for img_path in image_files: + img_name = os.path.basename(img_path) + + # Pattern 1: standard clip replace + mask_name = img_name.replace("_clip_", "_concatenated_ndwi_mask_clip_") + mask_path = os.path.join(self.data_dir, mask_name) + if os.path.exists(mask_path): + pairs.append((img_path, mask_path)) + continue + + # Pattern 2: clip replace without trailing clip + mask_name_alt1 = img_name.replace("_clip_", "_concatenated_ndwi_mask_") + mask_path_alt1 = os.path.join(self.data_dir, mask_name_alt1) + if os.path.exists(mask_path_alt1): + pairs.append((img_path, mask_path_alt1)) + continue + + # Pattern 3: regex base match + base_match = re.match(r'(.+)_\d+-of-\d+(_[^_]+)?\.tif$', img_name) + if base_match: + base_name = base_match.group(1) + suffix = base_match.group(2) if base_match.group(2) else "" + parts = img_name.replace('.tif', '').split('_') + tile_str = [p for p in parts if "-of-" in p] + if tile_str: + mask_name_alt2 = f"{base_name}_concatenated_ndwi_mask_{tile_str[0]}{suffix}.tif" + mask_path_alt2 = os.path.join(self.data_dir, mask_name_alt2) + if os.path.exists(mask_path_alt2): + pairs.append((img_path, mask_path_alt2)) + continue + print(f"Found {len(pairs)} 4-channel image-mask pairs in {self.data_dir}") + return pairs + + def __len__(self): + return len(self.image_mask_pairs) + + def __getitem__(self, idx): + img_path, mask_path = self.image_mask_pairs[idx] + + with rio.open(img_path) as src: + # Read all 4 bands (Blue, Green, Red, NIR) + image_data = src.read() # (4, H, W) + if image_data.shape[0] < 4: + pad_band = image_data[2:3] + image_data = np.concatenate([image_data, pad_band], axis=0) + elif image_data.shape[0] > 4: + image_data = image_data[:4] + + image_data = image_data.astype(np.float32) + image_data = np.nan_to_num(image_data, copy=False, nan=0.0, posinf=0.0, neginf=0.0) + + # DeepWaterMap style per-tile min-max normalization + img_min = np.min(image_data) + image_data = image_data - img_min + img_max = np.max(image_data) + if img_max > 0: + image_data = image_data / img_max + + with rio.open(mask_path) as src: + mask_data = src.read(1) + mask_data = (mask_data > 0).astype(np.float32) + + # Convert to torch tensors + img_tensor = torch.from_numpy(image_data) + mask_tensor = torch.from_numpy(mask_data).unsqueeze(0) + + # Resize to target input size + if img_tensor.shape[1:] != self.image_size: + img_tensor = F.interpolate(img_tensor.unsqueeze(0), size=self.image_size, mode='bilinear', align_corners=False).squeeze(0) + mask_tensor = F.interpolate(mask_tensor.unsqueeze(0), size=self.image_size, mode='nearest').squeeze(0) + + # Apply comprehensive online data augmentations during training + if self.is_train: + # 1. Random Horizontal Flip (50% probability) + if torch.rand(1).item() > 0.5: + img_tensor = torch.flip(img_tensor, dims=[2]) + mask_tensor = torch.flip(mask_tensor, dims=[2]) + + # 2. Random Vertical Flip (50% probability) + if torch.rand(1).item() > 0.5: + img_tensor = torch.flip(img_tensor, dims=[1]) + mask_tensor = torch.flip(mask_tensor, dims=[1]) + + # 3. Random 90°/180°/270° Rotation + k = torch.randint(0, 4, (1,)).item() + if k > 0: + img_tensor = torch.rot90(img_tensor, k=k, dims=[1, 2]) + mask_tensor = torch.rot90(mask_tensor, k=k, dims=[1, 2]) + + # 4. Random Brightness / Contrast Multiplicative Jitter (50% probability) + if torch.rand(1).item() > 0.5: + scale_factor = torch.empty(1).uniform_(0.85, 1.15).item() + img_tensor = torch.clamp(img_tensor * scale_factor, 0.0, 1.0) + + # 5. Random Additive Zero-Mean Gaussian Noise (50% probability) + if torch.rand(1).item() > 0.5: + noise_std = torch.empty(1).uniform_(0.005, 0.02).item() + noise = torch.randn_like(img_tensor) * noise_std + img_tensor = torch.clamp(img_tensor + noise, 0.0, 1.0) + + # 6. Random Scale & Crop (50% probability) + if torch.rand(1).item() > 0.5: + crop_ratio = torch.empty(1).uniform_(0.85, 1.0).item() + crop_h = int(self.image_size[0] * crop_ratio) + crop_w = int(self.image_size[1] * crop_ratio) + + top = torch.randint(0, self.image_size[0] - crop_h + 1, (1,)).item() + left = torch.randint(0, self.image_size[1] - crop_w + 1, (1,)).item() + + img_cropped = img_tensor[:, top:top+crop_h, left:left+crop_w] + mask_cropped = mask_tensor[:, top:top+crop_h, left:left+crop_w] + + img_tensor = F.interpolate(img_cropped.unsqueeze(0), size=self.image_size, mode='bilinear', align_corners=False).squeeze(0) + mask_tensor = F.interpolate(mask_cropped.unsqueeze(0), size=self.image_size, mode='nearest').squeeze(0) + + return img_tensor, mask_tensor diff --git a/training_pipeline/evaluate_unet.py b/training_pipeline/evaluate_unet.py new file mode 100644 index 0000000..c99bd35 --- /dev/null +++ b/training_pipeline/evaluate_unet.py @@ -0,0 +1,473 @@ +""" +U-Net Model Evaluation & Sept 4 / Sept 6 Coastline Analysis Script + +Evaluates trained U-Net / Attention U-Net models on: +1. September 4 and September 6 test tiles (test_data_4_6_sept folder). +2. Computes spatial RMSE metrics vs ground truth coastlines (ground_truth folder). +3. Computes regional RMSE breakdown across 5 coastal regions. +4. Computes pixel-level metrics (Pixel Accuracy, Precision, Recall, F1-Score, IoU) on validation set. + +Usage: python evaluate_unet.py +""" + +import os +import sys +import glob +import re +import numpy as np +import geopandas as gpd +import rasterio as rio +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, random_split +from torchvision import transforms +from PIL import Image +from tqdm import tqdm +import skimage.measure +from scipy import ndimage +from shapely.geometry import Point, MultiPoint, LineString, MultiLineString, box +import matplotlib.pyplot as plt + +# Dynamic linker fix & path resolution +script_dir = os.path.dirname(os.path.abspath(__file__)) +repo_root = os.path.abspath(os.path.join(script_dir, "..")) +if repo_root not in sys.path: + sys.path.append(repo_root) + +from load_config import load_config, get_training_config, get_augment_tiles_output_folder, get_model_save_path +from train_unet import UNet, AttentionUNet, ClassicUNet, SegmentationDataset + +UTM_ZONE_3N = 'EPSG:32603' +EPSILON = 2 ** -16 + +# ---------------------------- +# Line & Raster Smoothing +# ---------------------------- +def smooth_linestring(line, sigma=2.0, min_length=30.0): + """ + Applies 1D Gaussian rolling mean smoothing on LineString coordinates. + Preserves endpoints and filters out short noisy fragments. + """ + if line is None or line.is_empty or line.length < min_length: + return None + + coords = np.array(line.coords) + if len(coords) < 4: + return line + + xs, ys = coords[:, 0], coords[:, 1] + smoothed_xs = ndimage.gaussian_filter1d(xs, sigma=sigma, mode='nearest') + smoothed_ys = ndimage.gaussian_filter1d(ys, sigma=sigma, mode='nearest') + + smoothed_xs[0], smoothed_xs[-1] = xs[0], xs[-1] + smoothed_ys[0], smoothed_ys[-1] = ys[0], ys[-1] + + return LineString(np.column_stack([smoothed_xs, smoothed_ys])) + +def extract_model_coastline(model, image_path, transform, device, img_size=(256, 256)): + """ + Runs model inference on full resolution tile using sliding/tiled window or resized pass. + Extracts smoothed binary water mask & coastline LineStrings. + """ + model.eval() + with rio.open(image_path) as src: + data = src.read() # (4, H, W) + dataset_mask = src.dataset_mask() > 0 + h_orig, w_orig = data.shape[1], data.shape[2] + geo_transform = src.transform + + rgb = data[[2, 1, 0]] # (3, H, W) RGB + rgb = (np.clip(rgb.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) + rgb = np.transpose(rgb, (1, 2, 0)) + pil_img = Image.fromarray(rgb) + + input_t = transform(pil_img).unsqueeze(0).to(device) + + with torch.no_grad(): + raw_output = model(input_t) + if raw_output.shape[1] == 1: + prob_map = torch.sigmoid(raw_output).squeeze().cpu().numpy() + else: + prob_map = raw_output.squeeze().cpu().numpy() + + # Resize back to full raster resolution + prob_img = Image.fromarray((prob_map * 255).astype(np.uint8)).resize((w_orig, h_orig), Image.BILINEAR) + prob_map_full = np.array(prob_img, dtype=np.float32) / 255.0 + + # Probability smoothing & mask generation + smoothed_prob = ndimage.gaussian_filter(prob_map_full, sigma=1.5) + water_mask = (smoothed_prob > 0.5) & dataset_mask + + eroded_mask = ndimage.binary_erosion(dataset_mask, iterations=4) + contours = skimage.measure.find_contours(smoothed_prob, 0.5) + + lines = [] + for contour in contours: + rows = np.clip(np.round(contour[:, 0]).astype(int), 0, h_orig - 1) + cols = np.clip(np.round(contour[:, 1]).astype(int), 0, w_orig - 1) + + in_bounds = eroded_mask[rows, cols] + split_indices = np.where(~in_bounds)[0] + segments = np.split(contour, split_indices) + + for seg in segments: + seg_clean = seg[eroded_mask[np.clip(np.round(seg[:, 0]).astype(int), 0, h_orig - 1), + np.clip(np.round(seg[:, 1]).astype(int), 0, w_orig - 1)]] + if len(seg_clean) >= 2: + xs, ys = rio.transform.xy(geo_transform, seg_clean[:, 0], seg_clean[:, 1]) + raw_line = LineString(list(zip(xs, ys))) + smoothed_line = smooth_linestring(raw_line, sigma=2.0, min_length=30.0) + if smoothed_line is not None: + lines.append(smoothed_line) + + return lines, water_mask, geo_transform + +def extract_ndwi_coastline(image_path, transform): + """Extracts baseline NDWI threshold coastline.""" + with rio.open(image_path) as src: + green = src.read(2).astype(np.float32) + nir = src.read(4).astype(np.float32) + dataset_mask = src.dataset_mask() > 0 + h_orig, w_orig = green.shape + geo_transform = src.transform + + ndwi = (green - nir) / (green + nir + 1e-10) + smoothed_ndwi = ndimage.gaussian_filter(ndwi, sigma=1.5) + ndwi_mask = (smoothed_ndwi > 0.0) & dataset_mask + + eroded_mask = ndimage.binary_erosion(dataset_mask, iterations=4) + contours = skimage.measure.find_contours(smoothed_ndwi, 0.0) + + lines = [] + for contour in contours: + rows = np.clip(np.round(contour[:, 0]).astype(int), 0, h_orig - 1) + cols = np.clip(np.round(contour[:, 1]).astype(int), 0, w_orig - 1) + + in_bounds = eroded_mask[rows, cols] + split_indices = np.where(~in_bounds)[0] + segments = np.split(contour, split_indices) + + for seg in segments: + seg_clean = seg[eroded_mask[np.clip(np.round(seg[:, 0]).astype(int), 0, h_orig - 1), + np.clip(np.round(seg[:, 1]).astype(int), 0, w_orig - 1)]] + if len(seg_clean) >= 2: + xs, ys = rio.transform.xy(geo_transform, seg_clean[:, 0], seg_clean[:, 1]) + raw_line = LineString(list(zip(xs, ys))) + smoothed_line = smooth_linestring(raw_line, sigma=2.0, min_length=30.0) + if smoothed_line is not None: + lines.append(smoothed_line) + + return lines, ndwi_mask + +# ---------------------------- +# RMSE Computation Utilities +# ---------------------------- +def calc_rmse(errs): + errs = np.array(errs) + if len(errs) == 0: + return np.nan + return np.sqrt(np.square(errs).mean()) + +def find_distances(transects, fst, snd): + distances = [] + for transect in transects.itertuples(): + t_geom = transect.geometry + fst_pts = [p for p in fst.geoms if p.distance(t_geom) < EPSILON] + snd_pts = [p for p in snd.geoms if p.distance(t_geom) < EPSILON] + + if len(fst_pts) == 1 and len(snd_pts) == 1: + dist = fst_pts[0].distance(snd_pts[0]) + distances.append(dist) + return distances + +def compute_transect_rmse(transects, true_gdf, pred_lines_gdf, river_removal=True): + if pred_lines_gdf is None or len(pred_lines_gdf) == 0 or true_gdf is None or len(true_gdf) == 0: + return np.nan, [] + + if river_removal: + removal_ids = [17336, 17335, 17334, 17333, 17332] + transects = transects[~(transects['TransOrder'].isin(removal_ids))] + + transects = transects.to_crs(UTM_ZONE_3N) + true_gdf = true_gdf.to_crs(UTM_ZONE_3N) + pred_lines_gdf = pred_lines_gdf.to_crs(UTM_ZONE_3N) + + geom_true = true_gdf.unary_union.intersection(transects.unary_union) + geom_pred = pred_lines_gdf.unary_union.intersection(transects.unary_union) + + if geom_true.is_empty or geom_pred.is_empty: + return np.nan, [] + + if not hasattr(geom_true, 'geoms'): + geom_true = MultiPoint([geom_true]) if isinstance(geom_true, Point) else geom_true + if not hasattr(geom_pred, 'geoms'): + geom_pred = MultiPoint([geom_pred]) if isinstance(geom_pred, Point) else geom_pred + + dists = find_distances(transects, geom_true, geom_pred) + rmse_val = calc_rmse(dists) + return rmse_val, dists + +def compute_regional_rmse(transects, true_gdf, pred_lines_gdf): + regions = { + "Western Region (R1)": transects[transects['TransOrder'] >= 17443], + "Northern Region (R2)": transects[(transects['TransOrder'] < 17443) & (transects['TransOrder'] >= 17394)], + "Central Region (R3)": transects[(transects['TransOrder'] < 17394) & (transects['TransOrder'] >= 17370)], + "Town Region (R4)": transects[(transects['TransOrder'] < 17370) & (transects['TransOrder'] >= 17337)], + "East Region (R5)": transects[transects['TransOrder'] < 17337], + } + + res = {} + for r_name, r_transects in regions.items(): + val, _ = compute_transect_rmse(r_transects, true_gdf, pred_lines_gdf) + res[r_name] = val + return res + +# ---------------------------- +# Validation Dataset Evaluation +# ---------------------------- +def evaluate_validation_metrics(model, dataloader, device): + model.eval() + tp_total, fp_total, fn_total, tn_total = 0, 0, 0, 0 + pixel_correct, pixel_total = 0, 0 + + with torch.no_grad(): + for imgs, masks in tqdm(dataloader, desc="Evaluating Val Split"): + imgs, masks = imgs.to(device), masks.to(device) + outputs = model(imgs) + if outputs.shape[1] == 1: + outputs = torch.sigmoid(outputs) + preds = (outputs > 0.5).float() + + tp_total += ((preds == 1) & (masks == 1)).sum().item() + fp_total += ((preds == 1) & (masks == 0)).sum().item() + fn_total += ((preds == 0) & (masks == 1)).sum().item() + tn_total += ((preds == 0) & (masks == 0)).sum().item() + + pixel_correct += (preds == masks).sum().item() + pixel_total += masks.numel() + + acc = pixel_correct / pixel_total if pixel_total > 0 else 0 + prec = tp_total / (tp_total + fp_total) if (tp_total + fp_total) > 0 else 0 + rec = tp_total / (tp_total + fn_total) if (tp_total + fn_total) > 0 else 0 + f1 = (2 * prec * rec) / (prec + rec) if (prec + rec) > 0 else 0 + iou = tp_total / (tp_total + fp_total + fn_total) if (tp_total + fp_total + fn_total) > 0 else 0 + + return { + "pixel_accuracy": acc, + "precision": prec, + "recall": rec, + "f1_score": f1, + "iou": iou, + "tp": tp_total, + "fp": fp_total, + "fn": fn_total, + "tn": tn_total + } + +# ---------------------------- +# Plotting & Visualization +# ---------------------------- +def save_evaluation_plot(t_path, model_lines, ndwi_lines, planet_ref_gdf, usgs_gdf, hires_gt_gdf, out_plot_path): + with rio.open(t_path) as src: + rgb = src.read([3, 2, 1]) + bounds = src.bounds + + rgb_disp = np.zeros_like(rgb, dtype=np.float32) + for b in range(3): + band = rgb[b].astype(np.float32) + valid = band > 0 + if np.any(valid): + p2, p98 = np.percentile(band[valid], (2, 98)) + rgb_disp[b] = np.clip((band - p2) / (p98 - p2 + 1e-5), 0, 1) + + rgb_disp = np.transpose(rgb_disp, (1, 2, 0)) + + fig, ax = plt.subplots(figsize=(12, 12)) + ax.imshow(rgb_disp, extent=[bounds.left, bounds.right, bounds.bottom, bounds.top]) + + if planet_ref_gdf is not None: + planet_ref_gdf.to_crs(UTM_ZONE_3N).plot(ax=ax, color='orange', linewidth=2.5, label='Planet Labs Reference (9/9)') + if hires_gt_gdf is not None: + hires_gt_gdf.to_crs(UTM_ZONE_3N).plot(ax=ax, color='red', linewidth=2.0, label='Manual Hi-Res GT') + if usgs_gdf is not None: + usgs_gdf.to_crs(UTM_ZONE_3N).plot(ax=ax, color='green', linewidth=2.0, label='USGS Coastline') + + if model_lines: + model_gdf = gpd.GeoDataFrame(geometry=model_lines, crs=UTM_ZONE_3N) + model_gdf.plot(ax=ax, color='cyan', linewidth=1.8, label='Predicted Coastline (U-Net)') + + if ndwi_lines: + ndwi_gdf = gpd.GeoDataFrame(geometry=ndwi_lines, crs=UTM_ZONE_3N) + ndwi_gdf.plot(ax=ax, color='magenta', linewidth=1.5, label='NDWI Coastline') + + ax.set_title(f"Predicted Coastlines vs Ground Truth\n({os.path.basename(t_path)})", fontsize=14, fontweight='bold') + ax.set_xlim([bounds.left, bounds.right]) + ax.set_ylim([bounds.bottom, bounds.top]) + ax.legend(loc='upper right', fontsize=11) + plt.tight_layout() + + os.makedirs(os.path.dirname(out_plot_path), exist_ok=True) + plt.savefig(out_plot_path, dpi=300, bbox_inches='tight') + plt.close() + print(f"Saved evaluation comparison plot to: {out_plot_path}") + +# ---------------------------- +# Main Evaluation Function +# ---------------------------- +def main(): + config = load_config() + training_config = get_training_config(config) + + # Set device + device = training_config.get('device', 'auto') + if device == 'auto' or 'cuda' in device: + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {device}") + + # Paths to Ground Truth files (ground_truth directory) + gt_dir = os.path.join(repo_root, "ground_truth") + planet_ref_path = os.path.join(gt_dir, "9_9_16_PlanetCoastline.shp") + hires_gt_path = os.path.join(gt_dir, "2016_HiRes_Final_Coastline.shp") + + # Transects & USGS Coastlines + transects_path = os.path.join(repo_root, "USGS_Coastlines", "WestChukchi_exposed_STepr_rates", "WestChukchi_exposed_STepr_rates.shp") + + # Load Ground Truth shapefiles if available + planet_ref_gdf = gpd.read_file(planet_ref_path) if os.path.exists(planet_ref_path) else None + hires_gt_gdf = gpd.read_file(hires_gt_path) if os.path.exists(hires_gt_path) else None + usgs_gdf = gpd.read_file(transects_path) if os.path.exists(transects_path) else None + + if usgs_gdf is not None: + transects_gdf = usgs_gdf[usgs_gdf['BaselineID'] == 117] + else: + transects_gdf = None + + # Load Trained U-Net Model + model_paths = [ + os.path.join(repo_root, "output_models", "best_deep_attn_unet_8epochs.pth"), + os.path.join(repo_root, "output_models", "best_deep_unet_8epochs.pth"), + get_model_save_path(config) + ] + + model_path = None + for p in model_paths: + if os.path.exists(p): + model_path = p + break + + if model_path is None: + print("Warning: No pre-trained model weights found. Skipping spatial inference evaluation.") + else: + print(f"Loading model weights from: {model_path}") + state_dict = torch.load(model_path, map_location=device) + if isinstance(state_dict, dict) and 'model_state_dict' in state_dict: + state_dict = state_dict['model_state_dict'] + + is_attn = any("attn" in key for key in state_dict.keys()) + if is_attn: + model = AttentionUNet(n_channels=3, n_classes=1) + print("Loaded Attention U-Net architecture.") + else: + model = UNet(n_channels=3, n_classes=1) + print("Loaded Deep U-Net architecture.") + + model.load_state_dict(state_dict) + model.to(device) + + # Image preprocessing transform + image_size = training_config.get('image_size', [256, 256]) + transform = transforms.Compose([ + transforms.Resize(image_size), + transforms.ToTensor(), + ]) + + # Evaluate on September 4 and September 6 test tiles + test_dir = os.path.join(repo_root, "test_data_4_6_sept") + test_tiles = [ + os.path.join(test_dir, "sept_4", "files", "369619_2016-09-04_RE2_3A_Analytic_SR_clip.tif"), + os.path.join(test_dir, "sept_6", "files", "369619_2016-09-06_RE5_3A_Analytic_SR_clip.tif") + ] + + out_dir = os.path.join(repo_root, "inference_outputs") + os.makedirs(out_dir, exist_ok=True) + + for tile_path in test_tiles: + if not os.path.exists(tile_path): + print(f"Test tile not found: {tile_path}") + continue + + tile_name = os.path.basename(tile_path) + print(f"\n==================================================") + print(f"EVALUATING TEST TILE: {tile_name}") + print(f"==================================================") + + # Model inference & line extraction + model_lines, model_mask, geo_transform = extract_model_coastline(model, tile_path, transform, device) + ndwi_lines, ndwi_mask = extract_ndwi_coastline(tile_path, transform) + + model_gdf = gpd.GeoDataFrame(geometry=model_lines, crs=UTM_ZONE_3N) if model_lines else None + ndwi_gdf = gpd.GeoDataFrame(geometry=ndwi_lines, crs=UTM_ZONE_3N) if ndwi_lines else None + + # Save predicted shapefile + if model_gdf is not None: + shp_out = os.path.join(out_dir, f"{tile_name}_eval_predicted_coastline.shp") + model_gdf.to_file(shp_out) + print(f"Saved predicted coastline shapefile to: {shp_out}") + + # Compute Spatial RMSE Metrics + if transects_gdf is not None: + rmse_planet, _ = compute_transect_rmse(transects_gdf, planet_ref_gdf, model_gdf) + rmse_usgs, _ = compute_transect_rmse(transects_gdf, usgs_gdf, model_gdf) + rmse_hires, _ = compute_transect_rmse(transects_gdf, hires_gt_gdf, model_gdf) + + rmse_ndwi_planet, _ = compute_transect_rmse(transects_gdf, planet_ref_gdf, ndwi_gdf) + + print(f"\n[U-Net Model RMSE Results]") + print(f" RMSE vs Planet Labs Ref (9/9): {rmse_planet:.2f} m" if not np.isnan(rmse_planet) else " RMSE vs Planet Labs Ref: N/A") + print(f" RMSE vs USGS Coastlines: {rmse_usgs:.2f} m" if not np.isnan(rmse_usgs) else " RMSE vs USGS Coastlines: N/A") + print(f" RMSE vs Manual Hi-Res GT: {rmse_hires:.2f} m" if not np.isnan(rmse_hires) else " RMSE vs Manual Hi-Res GT: N/A") + + print(f"\n[NDWI Baseline RMSE Results]") + print(f" RMSE vs Planet Labs Ref (9/9): {rmse_ndwi_planet:.2f} m" if not np.isnan(rmse_ndwi_planet) else " RMSE vs Planet Labs Ref: N/A") + + # Regional Breakdown + if planet_ref_gdf is not None: + reg_unet = compute_regional_rmse(transects_gdf, planet_ref_gdf, model_gdf) + reg_ndwi = compute_regional_rmse(transects_gdf, planet_ref_gdf, ndwi_gdf) + + print(f"\n[Regional RMSE Breakdown vs Planet Labs Ref]") + for r_name in reg_unet.keys(): + u_v = f"{reg_unet[r_name]:.2f} m" if not np.isnan(reg_unet[r_name]) else "N/A" + n_v = f"{reg_ndwi[r_name]:.2f} m" if not np.isnan(reg_ndwi[r_name]) else "N/A" + print(f" - {r_name}: U-Net={u_v} | NDWI={n_v}") + + # Plot comparison map + plot_out_path = os.path.join(out_dir, f"{tile_name}_eval_comparison_plot.png") + save_evaluation_plot(tile_path, model_lines, ndwi_lines, planet_ref_gdf, usgs_gdf, hires_gt_gdf, plot_out_path) + + # Optional: Evaluate dataset split pixel metrics if augment_tiles folder exists + aug_data_dir = get_augment_tiles_output_folder(config) + if os.path.exists(aug_data_dir): + print(f"\n==========================================") + print(f"EVALUATING VALIDATION DATASET PIXEL METRICS") + print(f"==========================================") + val_dataset = SegmentationDataset(aug_data_dir, transform=transform) + if len(val_dataset) > 0: + train_split = training_config.get('train_split', 0.8) + total_sz = len(val_dataset) + train_sz = int(train_split * total_sz) + val_sz = total_sz - train_sz + + generator = torch.Generator().manual_seed(42) + _, val_set = random_split(val_dataset, [train_sz, val_sz], generator=generator) + val_loader = DataLoader(val_set, batch_size=16, num_workers=4, pin_memory=True) + + val_metrics = evaluate_validation_metrics(model, val_loader, device) + print(f"Pixel Accuracy: {val_metrics['pixel_accuracy']:.4%}") + print(f"Precision (PPV): {val_metrics['precision']:.4%}") + print(f"Recall (Sensitivity): {val_metrics['recall']:.4%}") + print(f"F1-Score (Dice Coeff): {val_metrics['f1_score']:.4%}") + print(f"Mean IoU (Jaccard): {val_metrics['iou']:.4%}") + +if __name__ == "__main__": + main() diff --git a/training_pipeline/predict.py b/training_pipeline/predict.py index 27886f7..0035129 100644 --- a/training_pipeline/predict.py +++ b/training_pipeline/predict.py @@ -13,40 +13,55 @@ # Add parent directory to path to import load_config and model sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from load_config import load_config, get_training_config, get_model_save_path -from train_unet import UNet +from train_unet import UNet, AttentionUNet, ClassicUNet + +def load_trained_model(model_path, device="cuda", model_type="auto"): + """ + Load a trained U-Net or Attention U-Net model. + """ + state_dict = torch.load(model_path, map_location=device) + if isinstance(state_dict, dict) and 'model_state_dict' in state_dict: + state_dict = state_dict['model_state_dict'] + + if model_type == "auto": + is_attn = any("attn" in key for key in state_dict.keys()) + if is_attn: + model = AttentionUNet(n_channels=3, n_classes=1) + else: + model = UNet(n_channels=3, n_classes=1) + elif model_type == "attention": + model = AttentionUNet(n_channels=3, n_classes=1) + elif model_type == "classic": + model = ClassicUNet(n_channels=3, n_classes=1) + else: + model = UNet(n_channels=3, n_classes=1) + + model.load_state_dict(state_dict) + model.to(device) + return model def predict_image(model, image_path, device="cuda", threshold=0.5): """ Predict mask for a single image. - - Args: - model: Trained U-Net model - image_path: Path to input image - device: Device to run inference on - threshold: Threshold for binary mask (default: 0.5) - - Returns: - numpy array: Predicted binary mask """ model.eval() - # Load and preprocess image image = Image.open(image_path).convert("RGB") original_size = image.size # (width, height) - # Use same transform as training transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), ]) - input_tensor = transform(image).unsqueeze(0).to(device) # add batch dimension + input_tensor = transform(image).unsqueeze(0).to(device) with torch.no_grad(): output = model(input_tensor) + if output.shape[1] == 1: + output = torch.sigmoid(output) pred_mask = (output.squeeze().cpu().numpy() > threshold).astype("uint8") - # Resize back to original size pred_mask_resized = Image.fromarray(pred_mask * 255).resize(original_size, Image.NEAREST) pred_mask_resized = np.array(pred_mask_resized) / 255 @@ -55,22 +70,10 @@ def predict_image(model, image_path, device="cuda", threshold=0.5): def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_path=None, resume=True): """ Predict masks for multiple images. - - Args: - model: Trained U-Net model - image_paths: List of paths to input images - device: Device to run inference on - threshold: Threshold for binary mask (default: 0.5) - checkpoint_path: Path to checkpoint file for resume capability - resume: Whether to resume from checkpoint if available - - Returns: - list: List of predicted binary masks """ predictions = [] processed_images = [] - # Try to load checkpoint if resume is enabled if resume and checkpoint_path: processed_images, predictions, metadata = load_checkpoint(checkpoint_path) if processed_images is None: @@ -79,7 +82,6 @@ def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_p else: print(f"Resuming from checkpoint: {len(processed_images)} images already processed") - # Process remaining images for i, image_path in enumerate(image_paths): if image_path in processed_images: print(f"Skipping already processed: {os.path.basename(image_path)}") @@ -91,7 +93,6 @@ def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_p predictions.append(pred_mask) processed_images.append(image_path) - # Save checkpoint every 10 images if checkpoint_path and (i + 1) % 10 == 0: metadata = { 'device': device, @@ -105,7 +106,6 @@ def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_p print(f"Error processing {image_path}: {e}") continue - # Final checkpoint save if checkpoint_path: metadata = { 'device': device, @@ -122,26 +122,18 @@ def predict_batch(model, image_paths, device="cuda", threshold=0.5, checkpoint_p def visualize_prediction(image_path, pred_mask, save_path=None): """ Visualize the original image and predicted mask side by side. - - Args: - image_path: Path to original image - pred_mask: Predicted mask array - save_path: Optional path to save the visualization """ fig, axes = plt.subplots(1, 3, figsize=(15, 5)) - # Original image original_image = Image.open(image_path).convert("RGB") axes[0].imshow(original_image) axes[0].set_title("Original Image") axes[0].axis('off') - # Predicted mask axes[1].imshow(pred_mask, cmap="gray") axes[1].set_title("Predicted Mask") axes[1].axis('off') - # Overlay overlay = np.array(original_image) overlay[pred_mask == 1] = [255, 0, 0] # Red overlay for predicted coastline axes[2].imshow(overlay) @@ -156,32 +148,7 @@ def visualize_prediction(image_path, pred_mask, save_path=None): plt.show() -def load_trained_model(model_path, device="cuda"): - """ - Load a trained U-Net model. - - Args: - model_path: Path to the saved model - device: Device to load the model on - - Returns: - Loaded U-Net model - """ - model = UNet(n_channels=3, n_classes=1) - model.load_state_dict(torch.load(model_path, map_location=device)) - model.to(device) - return model - def save_checkpoint(checkpoint_path, processed_images, predictions, metadata): - """ - Save prediction checkpoint with progress information. - - Args: - checkpoint_path (str): Path to save checkpoint file - processed_images (list): List of processed image paths - predictions (list): List of prediction results - metadata (dict): Additional metadata about the prediction run - """ checkpoint_data = { 'processed_images': processed_images, 'predictions': predictions, @@ -189,29 +156,16 @@ def save_checkpoint(checkpoint_path, processed_images, predictions, metadata): 'timestamp': datetime.now().isoformat(), 'total_processed': len(processed_images) } - with open(checkpoint_path, 'wb') as f: pickle.dump(checkpoint_data, f) - print(f"Checkpoint saved: {len(processed_images)} images processed") def load_checkpoint(checkpoint_path): - """ - Load prediction checkpoint to resume processing. - - Args: - checkpoint_path (str): Path to checkpoint file - - Returns: - tuple: (processed_images, predictions, metadata) or (None, None, None) if not found - """ if not os.path.exists(checkpoint_path): return None, None, None - try: with open(checkpoint_path, 'rb') as f: checkpoint_data = pickle.load(f) - print(f"Checkpoint loaded: {checkpoint_data['total_processed']} images already processed") return (checkpoint_data['processed_images'], checkpoint_data['predictions'], @@ -221,16 +175,6 @@ def load_checkpoint(checkpoint_path): return None, None, None def get_checkpoint_path(output_dir, batch_name="prediction_batch"): - """ - Generate checkpoint file path. - - Args: - output_dir (str): Directory to save checkpoint - batch_name (str): Name for the batch - - Returns: - str: Checkpoint file path - """ os.makedirs(output_dir, exist_ok=True) return os.path.join(output_dir, f"{batch_name}_checkpoint.pkl") @@ -248,17 +192,13 @@ def main(): args = parser.parse_args() - # Load configuration config = load_config() - training_config = get_training_config(config) - # Set device if args.device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" else: device = args.device - # Get model path if args.model: model_path = args.model else: @@ -267,16 +207,13 @@ def main(): print(f"Loading model from: {model_path}") print(f"Using device: {device}") - # Load model if not os.path.exists(model_path): print(f"Model file not found: {model_path}") - print("Please train the model first or provide a valid model path with --model") return model = load_trained_model(model_path, device) print("Model loaded successfully") - # Process single image if args.image: if not os.path.exists(args.image): print(f"Image file not found: {args.image}") @@ -284,24 +221,19 @@ def main(): print(f"Predicting mask for: {args.image}") pred_mask = predict_image(model, args.image, device, args.threshold) - - # Visualize result visualize_prediction(args.image, pred_mask, args.save) - # Save mask if save path provided if args.save and not args.save.endswith(('.png', '.jpg', '.jpeg')): mask_save_path = args.save.replace('.png', '_mask.png') Image.fromarray(pred_mask * 255).save(mask_save_path) print(f"Mask saved to: {mask_save_path}") - # Process multiple images elif args.images: valid_images = [img for img in args.images if os.path.exists(img)] if not valid_images: print("No valid image files found") return - # Set up checkpoint path if checkpoint directory is provided checkpoint_path = None resume = not args.no_resume if args.checkpoint_dir: @@ -312,7 +244,6 @@ def main(): predictions = predict_batch(model, valid_images, device, args.threshold, checkpoint_path, resume) - # Save predictions for i, (image_path, pred_mask) in enumerate(zip(valid_images, predictions)): base_name = os.path.splitext(os.path.basename(image_path))[0] mask_save_path = f"{base_name}_predicted_mask.png" diff --git a/training_pipeline/run_inference_sept_4_6.py b/training_pipeline/run_inference_sept_4_6.py deleted file mode 100644 index 3c02b73..0000000 --- a/training_pipeline/run_inference_sept_4_6.py +++ /dev/null @@ -1,698 +0,0 @@ -import os -import sys -import numpy as np -import geopandas as gpd -import rasterio as rio -import torch -import torch.nn as nn -import torch.nn.functional as F -from torchvision import transforms -from PIL import Image -import skimage.measure -from shapely.geometry import Point, MultiPoint, LineString, MultiLineString, box -import matplotlib.pyplot as plt -import scipy.ndimage as ndimage -import matplotlib.colors as mcolors -from matplotlib.patches import Patch - -# Dynamic path resolution to find CoastlineExtraction root and training_pipeline -script_dir = os.path.dirname(os.path.abspath(__file__)) -candidate_roots = [ - os.path.abspath(os.path.join(script_dir, "..")), - os.path.abspath(os.path.join(script_dir, "..", "CoastlineExtraction")), - script_dir, -] - -repo_root = None -for candidate in candidate_roots: - if os.path.exists(os.path.join(candidate, "training_pipeline")) or os.path.exists(os.path.join(candidate, "test_data_4_6_sept")): - repo_root = candidate - break - -if repo_root is None: - repo_root = candidate_roots[0] - -sys_paths = [ - os.path.join(repo_root, "training_pipeline"), - script_dir, - repo_root, -] -for p in sys_paths: - if p not in sys.path and os.path.exists(p): - sys.path.append(p) - -from train_and_eval_pipeline import UNet, AttentionUNet - -UTM_ZONE_3N = 'EPSG:32603' - -# --------------------------------------------- -# DeepWaterMap PyTorch Architecture Definition -# --------------------------------------------- -class ConvBlockDWM(nn.Module): - def __init__(self, in_c, out_c, k_size, stride=1, use_relu=True): - super().__init__() - padding = k_size // 2 - self.conv = nn.Conv2d(in_c, out_c, kernel_size=k_size, stride=stride, padding=padding, bias=False) - self.bn = nn.BatchNorm2d(out_c, eps=1e-3, momentum=0.01) - self.use_relu = use_relu - - def forward(self, x): - x = self.conv(x) - x = self.bn(x) - if self.use_relu: - x = F.relu(x) - return x - -class DownscalingUnitDWM(nn.Module): - def __init__(self, in_c, out_c): - super().__init__() - self.c1 = ConvBlockDWM(in_c, out_c, k_size=5, stride=2, use_relu=True) - self.c2 = ConvBlockDWM(out_c, out_c, k_size=3, stride=1, use_relu=True) - - def forward(self, x): - x1 = self.c1(x) - x2 = self.c2(x1) - return x1 + x2 - -class UpscalingUnitDWM(nn.Module): - def __init__(self, in_c, out_c): - super().__init__() - self.pixel_shuffle = nn.PixelShuffle(2) - self.c1 = ConvBlockDWM(in_c // 4, out_c, k_size=3, stride=1, use_relu=True) - self.c2 = ConvBlockDWM(out_c, out_c, k_size=3, stride=1, use_relu=True) - - def forward(self, x): - x = self.pixel_shuffle(x) - x1 = self.c1(x) - x2 = self.c2(x1) - return x1 + x2 - -class BottleneckUnitDWM(nn.Module): - def __init__(self, c): - super().__init__() - self.c1 = ConvBlockDWM(c, c, k_size=3, stride=1, use_relu=True) - self.c2 = ConvBlockDWM(c, c, k_size=3, stride=1, use_relu=True) - - def forward(self, x): - x1 = self.c1(x) - x2 = self.c2(x1) - return x1 + x2 - -class DeepWaterMapPyTorch(nn.Module): - def __init__(self): - super().__init__() - self.first_layer = ConvBlockDWM(6, 4, k_size=1, stride=1, use_relu=False) - self.down1 = DownscalingUnitDWM(4, 16) - self.down2 = DownscalingUnitDWM(16, 64) - self.down3 = DownscalingUnitDWM(64, 256) - self.down4 = DownscalingUnitDWM(256, 1024) - - self.bottleneck = BottleneckUnitDWM(1024) - - self.up1 = UpscalingUnitDWM(1024, 256) - self.up2 = UpscalingUnitDWM(256, 64) - self.up3 = UpscalingUnitDWM(64, 16) - self.up4 = UpscalingUnitDWM(16, 4) - - self.last_layer = ConvBlockDWM(4, 1, k_size=1, stride=1, use_relu=False) - - def forward(self, x): - skips = [] - x0 = self.first_layer(x) - skips.append(x0) - - x1 = self.down1(x0) - skips.append(x1) - - x2 = self.down2(x1) - skips.append(x2) - - x3 = self.down3(x2) - skips.append(x3) - - x4 = self.down4(x3) - skips.append(x4) - - b = self.bottleneck(x4) - - d1 = b + skips.pop() - u1 = self.up1(d1) - - d2 = u1 + skips.pop() - u2 = self.up2(d2) - - d3 = u2 + skips.pop() - u3 = self.up3(d3) - - d4 = u3 + skips.pop() - u4 = self.up4(d4) - - d_last = u4 + skips.pop() - out = self.last_layer(d_last) - return torch.sigmoid(out) - -# ---------------------------- -# Helper & Inference Functions -# ---------------------------- -def smooth_linestring(line, sigma=2.0, min_length=30.0): - """ - Smooths a Shapely LineString using a 1D Gaussian rolling mean filter along X and Y coordinates. - Filters out short noisy fragments shorter than min_length (meters). - """ - if line is None or line.is_empty or line.length < min_length: - return None - - coords = np.array(line.coords) - if len(coords) < 4: - return line - - xs, ys = coords[:, 0], coords[:, 1] - smoothed_xs = ndimage.gaussian_filter1d(xs, sigma=sigma, mode='nearest') - smoothed_ys = ndimage.gaussian_filter1d(ys, sigma=sigma, mode='nearest') - - # Preserve exact line endpoints to avoid shrinking - smoothed_xs[0], smoothed_xs[-1] = xs[0], xs[-1] - smoothed_ys[0], smoothed_ys[-1] = ys[0], ys[-1] - - return LineString(np.column_stack([smoothed_xs, smoothed_ys])) - -def stretch_rgb_image(image_data, dataset_mask=None): - """ - Applies a 2%-98% percentile contrast stretch on valid non-zero satellite pixels - so the RGB visualization matches bright QGIS rendering. - """ - scaled = np.zeros_like(image_data, dtype=np.uint8) - for c in range(3): - channel = image_data[:, :, c].astype(np.float32) - valid = channel[dataset_mask] if dataset_mask is not None else channel[channel > 0] - if len(valid) > 0: - p2, p98 = np.percentile(valid, (2, 98)) - if p98 > p2: - channel = np.clip((channel - p2) / (p98 - p2) * 255.0, 0, 255) - scaled[:, :, c] = channel.astype(np.uint8) - return scaled - -def extract_model_coastline(model, image_path, transform, device): - model.eval() - with rio.open(image_path) as src: - image_data = src.read([3, 2, 1]) - dataset_mask = src.dataset_mask() > 0 - h_orig, w_orig = image_data.shape[1], image_data.shape[2] - - scaled_data = (np.clip(image_data.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) - scaled_data = np.transpose(scaled_data, (1, 2, 0)) - - image = Image.fromarray(scaled_data) - transform_resize = transforms.Compose([ - transforms.Resize((256, 256)), - transforms.ToTensor(), - ]) - img_tensor = transform_resize(image).unsqueeze(0).to(device) - - with torch.no_grad(): - output = torch.sigmoid(model(img_tensor)).squeeze().cpu().numpy() - - output_resized = Image.fromarray((output * 255).astype(np.uint8)).resize((w_orig, h_orig), Image.BILINEAR) - prob_map = np.array(output_resized, dtype=np.float32) / 255.0 - - # Smooth spatial probability map with Gaussian filter before contouring - smoothed_prob = ndimage.gaussian_filter(prob_map, sigma=1.5) - pred_mask_np = (smoothed_prob > 0.5) & dataset_mask - - eroded_mask = ndimage.binary_erosion(dataset_mask, iterations=4) - contours = skimage.measure.find_contours(smoothed_prob, 0.5) - - lines = [] - for contour in contours: - rows = np.clip(np.round(contour[:, 0]).astype(int), 0, h_orig - 1) - cols = np.clip(np.round(contour[:, 1]).astype(int), 0, w_orig - 1) - - in_bounds = eroded_mask[rows, cols] - split_indices = np.where(~in_bounds)[0] - segments = np.split(contour, split_indices) - for seg in segments: - seg_clean = seg[eroded_mask[np.clip(np.round(seg[:, 0]).astype(int), 0, h_orig - 1), - np.clip(np.round(seg[:, 1]).astype(int), 0, w_orig - 1)]] - if len(seg_clean) >= 2: - xs, ys = rio.transform.xy(transform, seg_clean[:, 0], seg_clean[:, 1]) - raw_line = LineString(list(zip(xs, ys))) - smoothed_line = smooth_linestring(raw_line, sigma=2.5, min_length=40.0) - if smoothed_line is not None: - lines.append(smoothed_line) - return lines, pred_mask_np - -def extract_ndwi_coastline(image_path, transform): - """ - Computes NDWI = (Green - NIR) / (Green + NIR) and extracts the resulting smoothed coastline contour. - """ - with rio.open(image_path) as src: - green = src.read(2).astype(np.float32) # Band 2: Green - nir = src.read(src.count).astype(np.float32) # Band 4: NIR - dataset_mask = src.dataset_mask() > 0 - h_orig, w_orig = green.shape - - denom = green + nir - denom[denom == 0] = 1e-5 - ndwi = (green - nir) / denom - - # Spatial Gaussian smoothing on NDWI array to reduce pixel noise - ndwi_smoothed = ndimage.gaussian_filter(ndwi, sigma=1.5) - ndwi_mask = (ndwi_smoothed > 0.0) & dataset_mask - eroded_mask = ndimage.binary_erosion(dataset_mask, iterations=4) - contours = skimage.measure.find_contours(ndwi_smoothed, 0.0) - - lines = [] - for contour in contours: - rows = np.clip(np.round(contour[:, 0]).astype(int), 0, h_orig - 1) - cols = np.clip(np.round(contour[:, 1]).astype(int), 0, w_orig - 1) - - in_bounds = eroded_mask[rows, cols] - split_indices = np.where(~in_bounds)[0] - segments = np.split(contour, split_indices) - for seg in segments: - seg_clean = seg[eroded_mask[np.clip(np.round(seg[:, 0]).astype(int), 0, h_orig - 1), - np.clip(np.round(seg[:, 1]).astype(int), 0, w_orig - 1)]] - if len(seg_clean) >= 2: - xs, ys = rio.transform.xy(transform, seg_clean[:, 0], seg_clean[:, 1]) - raw_line = LineString(list(zip(xs, ys))) - smoothed_line = smooth_linestring(raw_line, sigma=2.5, min_length=40.0) - if smoothed_line is not None: - lines.append(smoothed_line) - return lines, ndwi_mask - -def extract_dwm_coastline(dwm_model, image_path, transform, device): - """ - Runs DeepWaterMap model inference to extract surface water coastline contour. - """ - dwm_model.eval() - with rio.open(image_path) as src: - data = src.read() # (4, H, W) - dataset_mask = src.dataset_mask() > 0 - h_orig, w_orig = data.shape[1], data.shape[2] - - pad_r = (32 - h_orig % 32) % 32 - pad_c = (32 - w_orig % 32) % 32 - - data_6b = np.zeros((6, h_orig, w_orig), dtype=np.float32) - data_6b[0] = data[0] # Blue - data_6b[1] = data[1] # Green - data_6b[2] = data[2] # Red - data_6b[3] = data[3] # NIR - data_6b[4] = data[3] # SWIR1 approx - data_6b[5] = data[3] # SWIR2 approx - - data_padded = np.pad(data_6b, ((0, 0), (0, pad_r), (0, pad_c)), 'reflect') - data_padded = np.nan_to_num(data_padded, copy=False, nan=0.0, posinf=0.0, neginf=0.0) - min_v = np.min(data_padded) - max_v = np.maximum(np.max(data_padded), 1.0) - data_padded = (data_padded - min_v) / max_v - - inp_t = torch.from_numpy(data_padded).unsqueeze(0).to(device) - with torch.no_grad(): - raw_pred = dwm_model(inp_t).squeeze().cpu().numpy() - - if pad_r > 0: raw_pred = raw_pred[:-pad_r, :] - if pad_c > 0: raw_pred = raw_pred[:, :-pad_c] - - soft_pred = 1.0 / (1.0 + np.exp(-(16.0 * (raw_pred - 0.5)))) - soft_pred = np.clip(soft_pred, 0, 1) - - smoothed_dwm = ndimage.gaussian_filter(soft_pred, sigma=1.0) - dwm_mask = (smoothed_dwm > 0.5) & dataset_mask - - eroded_mask = ndimage.binary_erosion(dataset_mask, iterations=4) - contours = skimage.measure.find_contours(smoothed_dwm, 0.5) - - lines = [] - for contour in contours: - rows = np.clip(np.round(contour[:, 0]).astype(int), 0, h_orig - 1) - cols = np.clip(np.round(contour[:, 1]).astype(int), 0, w_orig - 1) - - in_bounds = eroded_mask[rows, cols] - split_indices = np.where(~in_bounds)[0] - segments = np.split(contour, split_indices) - for seg in segments: - seg_clean = seg[eroded_mask[np.clip(np.round(seg[:, 0]).astype(int), 0, h_orig - 1), - np.clip(np.round(seg[:, 1]).astype(int), 0, w_orig - 1)]] - if len(seg_clean) >= 2: - xs, ys = rio.transform.xy(transform, seg_clean[:, 0], seg_clean[:, 1]) - raw_line = LineString(list(zip(xs, ys))) - smoothed_line = smooth_linestring(raw_line, sigma=2.0, min_length=15.0) - if smoothed_line is not None: - lines.append(smoothed_line) - return lines, dwm_mask - -def plot_water_land_prediction(t_path, pred_mask_np, pred_lines, plot_out_path): - fig, axes = plt.subplots(1, 3, figsize=(18, 6)) - tile_name = os.path.basename(t_path) - - with rio.open(t_path) as src: - image_data = src.read([3, 2, 1]) - image_data = np.transpose(image_data, (1, 2, 0)) - dataset_mask = src.dataset_mask() > 0 - tile_bounds = src.bounds - extent = [tile_bounds.left, tile_bounds.right, tile_bounds.bottom, tile_bounds.top] - - rgb_stretched = stretch_rgb_image(image_data, dataset_mask) - - axes[0].imshow(rgb_stretched, extent=extent) - axes[0].set_title("RGB Satellite Image (QGIS Stretch)", fontsize=12, fontweight="bold") - axes[0].axis("off") - - class_map = np.zeros(pred_mask_np.shape, dtype=np.uint8) - class_map[pred_mask_np] = 1 - class_map[~dataset_mask] = 2 - - cmap_water_land = mcolors.ListedColormap(['#8B5A2B', '#1E88E5', '#FFFFFF']) - axes[1].imshow(class_map, extent=extent, cmap=cmap_water_land) - axes[1].set_title("Binary Classification (Water vs Land)", fontsize=12, fontweight="bold") - axes[1].axis("off") - - legend_elements = [ - Patch(facecolor='#1E88E5', label='Water (Predicted)'), - Patch(facecolor='#8B5A2B', label='Land (Predicted)'), - Patch(facecolor='#FFFFFF', edgecolor='gray', label='NoData (Background)') - ] - axes[1].legend(handles=legend_elements, loc="upper right") - - axes[2].imshow(rgb_stretched, extent=extent) - water_overlay = np.zeros((*pred_mask_np.shape, 4), dtype=np.float32) - water_overlay[pred_mask_np] = [0.12, 0.53, 0.90, 0.45] - axes[2].imshow(water_overlay, extent=extent) - - for idx, line in enumerate(pred_lines): - lbl = "Extracted Coastline (U-Net)" if idx == 0 else "" - axes[2].plot(*line.xy, color="cyan", linewidth=1.2, linestyle="-", label=lbl) - if pred_lines: - axes[2].legend(loc="upper right") - - axes[2].set_title("Water Mask & Clean Coastline Overlay", fontsize=12, fontweight="bold") - axes[2].axis("off") - - fig.suptitle(f"Model Binary Water/Land Prediction Analysis - {tile_name}", fontsize=14, fontweight="bold", y=0.98) - plt.tight_layout() - plt.savefig(plot_out_path, dpi=300, bbox_inches="tight") - plt.close() - print(f"Saved binary water/land visualization plot to: {plot_out_path}") - -def calculate_rmse_on_transects(predicted_lines, trans_deering, ref_distances, usgs_distances, hires_distances): - empty_regional = {1: float('nan'), 2: float('nan'), 3: float('nan'), 4: float('nan'), 5: float('nan')} - if not predicted_lines: - return float('nan'), float('nan'), float('nan'), empty_regional - - combined_lines = MultiLineString(predicted_lines) - errors_planet = [] - errors_usgs = [] - errors_hires = [] - regional_errors = {1: [], 2: [], 3: [], 4: [], 5: []} - - for idx, row in trans_deering.iterrows(): - oid = int(row['TransOrder']) - t_geom = row.geometry - - region_id = 5 - if oid >= 17443: - region_id = 1 # Western Region - elif oid >= 17394: - region_id = 2 # Northern Region - elif oid >= 17370: - region_id = 3 # Central Region - elif oid >= 17337: - region_id = 4 # Town Region - - pt_int = combined_lines.intersection(t_geom) - dist = None - if not pt_int.is_empty: - if isinstance(pt_int, Point): - dist = t_geom.project(pt_int) - elif isinstance(pt_int, MultiPoint): - dist = min([t_geom.project(pt) for pt in pt_int.geoms]) - elif hasattr(pt_int, 'geoms'): - pts = [pt for pt in pt_int.geoms if isinstance(pt, Point)] - if pts: - dist = min([t_geom.project(pt) for pt in pts]) - - if dist is not None: - if oid in ref_distances: - err_p = dist - ref_distances[oid] - errors_planet.append(err_p) - regional_errors[region_id].append(err_p) - if oid in usgs_distances: - err_u = dist - usgs_distances[oid] - errors_usgs.append(err_u) - if oid in hires_distances: - err_h = dist - hires_distances[oid] - errors_hires.append(err_h) - - rmse_planet = np.sqrt(np.mean(np.square(errors_planet))) if errors_planet else float('nan') - rmse_usgs = np.sqrt(np.mean(np.square(errors_usgs))) if errors_usgs else float('nan') - rmse_hires = np.sqrt(np.mean(np.square(errors_hires))) if errors_hires else float('nan') - - regional_rmse = {} - for r_id, errs in regional_errors.items(): - regional_rmse[r_id] = np.sqrt(np.mean(np.square(errs))) if errs else float('nan') - - return rmse_planet, rmse_usgs, rmse_hires, regional_rmse - -def plot_predictions_comparison(t_path, u_lines, ndwi_lines, dwm_lines, planet_union, hires_union, usgs_union, transform, plot_out_path): - fig, ax = plt.subplots(figsize=(10, 10)) - ax.set_title(f"Predicted Coastlines vs. Ground Truth\n(Tile: {os.path.basename(t_path)})", fontsize=14, fontweight="bold") - - with rio.open(t_path) as src: - image_data = src.read([3, 2, 1]) - image_data = np.transpose(image_data, (1, 2, 0)) - dataset_mask = src.dataset_mask() > 0 - tile_bounds = src.bounds - extent = [tile_bounds.left, tile_bounds.right, tile_bounds.bottom, tile_bounds.top] - - rgb_stretched = stretch_rgb_image(image_data, dataset_mask) - ax.imshow(rgb_stretched, extent=extent, alpha=0.95) - tile_box = box(*tile_bounds) - - p_cropped = planet_union.intersection(tile_box) - h_cropped = hires_union.intersection(tile_box) - u_cropped = usgs_union.intersection(tile_box) - - # Plot Ground Truths - def plot_geom(geom, color, linewidth, label, linestyle="-"): - if geom.is_empty: - return - if isinstance(geom, LineString): - ax.plot(*geom.xy, color=color, linewidth=linewidth, label=label, linestyle=linestyle) - elif isinstance(geom, MultiLineString): - for line in geom.geoms: - ax.plot(*line.xy, color=color, linewidth=linewidth, label=label, linestyle=linestyle) - elif hasattr(geom, "geoms"): - for sub_geom in geom.geoms: - plot_geom(sub_geom, color, linewidth, label, linestyle) - - plot_geom(p_cropped, color="orange", linewidth=2.0, label="Planet Labs Reference", linestyle="-") - plot_geom(h_cropped, color="red", linewidth=2.0, label="Manual Hi-Res GT", linestyle="-") - plot_geom(u_cropped, color="green", linewidth=2.0, label="USGS Coastline", linestyle="-") - - # Plot predicted coastlines (solid thin lines) - for idx, line in enumerate(u_lines): - lbl = "Predicted Coastline (U-Net)" if idx == 0 else "" - ax.plot(*line.xy, color="cyan", linewidth=1.2, linestyle="-", label=lbl) - - for idx, line in enumerate(ndwi_lines): - lbl = "NDWI Coastline" if idx == 0 else "" - ax.plot(*line.xy, color="magenta", linewidth=1.2, linestyle="-", label=lbl) - - for idx, line in enumerate(dwm_lines): - lbl = "DeepWaterMap Coastline" if idx == 0 else "" - ax.plot(*line.xy, color="yellow", linewidth=1.2, linestyle="-", label=lbl) - - handles, labels = ax.get_legend_handles_labels() - by_label = {} - for h, l in zip(handles, labels): - if l: - by_label[l] = h - ax.legend(by_label.values(), by_label.keys(), loc="upper right") - - plt.tight_layout() - plt.savefig(plot_out_path, dpi=300) - plt.close() - print(f"Saved visualization plot to: {plot_out_path}") - -def main(): - print("==================================================") - print("RUNNING MODEL INFERENCE (U-NET, NDWI & DEEPWATERMAP)") - print("==================================================") - - device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") - print(f"Using device: {device}") - - output_dir = os.path.join(repo_root, "output_models") - attn_model_path = os.path.join(output_dir, "best_deep_attn_unet_8epochs.pth") - unet_model_path = os.path.join(output_dir, "best_deep_unet_8epochs.pth") - legacy_model_path = os.path.join(output_dir, "best_unet_8epochs.pth") - dwm_model_path = os.path.join(output_dir, "deepwatermap_pytorch.pth") - - if os.path.exists(attn_model_path): - model_path = attn_model_path - model = AttentionUNet(n_channels=3, n_classes=1) - model_type = "Attention U-Net" - elif os.path.exists(unet_model_path): - model_path = unet_model_path - model = UNet(n_channels=3, n_classes=1) - model_type = "Standard U-Net" - elif os.path.exists(legacy_model_path): - model_path = legacy_model_path - model = UNet(n_channels=3, n_classes=1) - model_type = "Legacy Standard U-Net" - else: - print(f"Error: No model weights found at {attn_model_path} or {unet_model_path}.") - sys.exit(1) - - print(f"Loading [{model_type}] weights from: {model_path}") - model.load_state_dict(torch.load(model_path, map_location=device)) - model.to(device) - model.eval() - - dwm_model = None - if os.path.exists(dwm_model_path): - print(f"Loading [DeepWaterMap] weights from: {dwm_model_path}") - dwm_model = DeepWaterMapPyTorch().to(device) - dwm_model.load_state_dict(torch.load(dwm_model_path, map_location=device)) - dwm_model.eval() - else: - print(f"Warning: DeepWaterMap weights not found at {dwm_model_path}.") - - p_transects = os.path.join(repo_root, "USGS_Coastlines", "WestChukchi_exposed_STepr_rates", "WestChukchi_exposed_STepr_rates.shp") - p_usgs = os.path.join(repo_root, "USGS_Coastlines", "Deering_shorelines_2016.shp") - p_hires = os.path.join(repo_root, "ground_truth", "2016_HiRes_Final_Coastline.shp") - - planet_candidates = [ - os.path.join(repo_root, "existing_data", "DigitizedCoastlines", "PlanetCoastline_gt", "09_09_2016", "9_9_16_PlanetCoastline.shp"), - os.path.join(repo_root, "..", "existing_data", "DigitizedCoastlines", "PlanetCoastline_gt", "09_09_2016", "9_9_16_PlanetCoastline.shp") - ] - p_planet = next((p for p in planet_candidates if os.path.exists(p)), planet_candidates[0]) - - print("Loading transects and reference coastlines...") - trans_deering = gpd.read_file(p_transects).to_crs(UTM_ZONE_3N) - trans_deering = trans_deering[trans_deering['BaselineID'] == 117].sort_values('TransOrder') - usgs = gpd.read_file(p_usgs).to_crs(UTM_ZONE_3N) - hires = gpd.read_file(p_hires).to_crs(UTM_ZONE_3N) - planet = gpd.read_file(p_planet).to_crs(UTM_ZONE_3N) - - usgs_union = usgs.union_all() if hasattr(usgs, 'union_all') else usgs.unary_union - hires_union = hires.union_all() if hasattr(hires, 'union_all') else hires.unary_union - planet_union = planet.union_all() if hasattr(planet, 'union_all') else planet.unary_union - - planet_distances = {} - usgs_distances = {} - hires_distances = {} - for idx, row in trans_deering.iterrows(): - oid = int(row['TransOrder']) - t_geom = row.geometry - - p_int = t_geom.intersection(planet_union) - if not p_int.is_empty: - if isinstance(p_int, Point): - planet_distances[oid] = t_geom.project(p_int) - elif isinstance(p_int, MultiPoint): - planet_distances[oid] = min([t_geom.project(pt) for pt in p_int.geoms]) - elif hasattr(p_int, 'geoms'): - pts = [pt for pt in p_int.geoms if isinstance(pt, Point)] - if pts: - planet_distances[oid] = min([t_geom.project(pt) for pt in pts]) - - u_int = t_geom.intersection(usgs_union) - if not u_int.is_empty: - if isinstance(u_int, Point): - usgs_distances[oid] = t_geom.project(u_int) - elif isinstance(u_int, MultiPoint): - usgs_distances[oid] = min([t_geom.project(pt) for pt in u_int.geoms]) - elif hasattr(u_int, 'geoms'): - pts = [pt for pt in u_int.geoms if isinstance(pt, Point)] - if pts: - usgs_distances[oid] = min([t_geom.project(pt) for pt in pts]) - - h_int = t_geom.intersection(hires_union) - if not h_int.is_empty: - if isinstance(h_int, Point): - hires_distances[oid] = t_geom.project(h_int) - elif isinstance(h_int, MultiPoint): - hires_distances[oid] = min([t_geom.project(pt) for pt in h_int.geoms]) - elif hasattr(h_int, 'geoms'): - pts = [pt for pt in h_int.geoms if isinstance(pt, Point)] - if pts: - hires_distances[oid] = min([t_geom.project(pt) for pt in pts]) - - test_tiles = [ - os.path.join(repo_root, "test_data_4_6_sept", "sept_4", "files", "369619_2016-09-04_RE2_3A_Analytic_SR_clip.tif"), - os.path.join(repo_root, "test_data_4_6_sept", "sept_6", "files", "369619_2016-09-06_RE5_3A_Analytic_SR_clip.tif") - ] - - vis_dir = os.path.join(repo_root, "inference_outputs") - os.makedirs(vis_dir, exist_ok=True) - - region_names = { - 1: "Western Region", - 2: "Northern Region", - 3: "Central Region", - 4: "Town Region", - 5: "East Region" - } - - for t_path in test_tiles: - print(f"\nProcessing test tile: {os.path.basename(t_path)}") - if not os.path.exists(t_path): - print(f"Error: Tile {t_path} not found.") - continue - - with rio.open(t_path) as src: - transform = src.transform - - pred_lines, pred_mask_np = extract_model_coastline(model, t_path, transform, device) - ndwi_lines, ndwi_mask_np = extract_ndwi_coastline(t_path, transform) - - dwm_lines, dwm_mask_np = [], np.zeros_like(pred_mask_np) - if dwm_model is not None: - dwm_lines, dwm_mask_np = extract_dwm_coastline(dwm_model, t_path, transform, device) - - tile_name = os.path.splitext(os.path.basename(t_path))[0] - shp_out_path = os.path.join(vis_dir, f"{tile_name}_model_predicted_coastline.shp") - if pred_lines: - gdf_pred = gpd.GeoDataFrame(geometry=pred_lines, crs=UTM_ZONE_3N) - gdf_pred.to_file(shp_out_path) - print(f"Saved predicted coastline shapefile to: {shp_out_path}") - - # Calculate RMSE scores - rmse_p, rmse_u, rmse_h, regional_rmse = calculate_rmse_on_transects(pred_lines, trans_deering, planet_distances, usgs_distances, hires_distances) - ndwi_p, ndwi_u, ndwi_h, ndwi_regional = calculate_rmse_on_transects(ndwi_lines, trans_deering, planet_distances, usgs_distances, hires_distances) - dwm_p, dwm_u, dwm_h, dwm_regional = calculate_rmse_on_transects(dwm_lines, trans_deering, planet_distances, usgs_distances, hires_distances) - - print(f"\nRESULTS FOR TILE: {os.path.basename(t_path)}") - print(" [U-Net Model]") - print(f" RMSE vs Planet Labs Ref: {rmse_p:.2f} m") - print(f" RMSE vs USGS Coastlines: {rmse_u:.2f} m") - print(f" RMSE vs Manual Hi-Res GT: {rmse_h:.2f} m") - - print(" [NDWI Threshold]") - print(f" RMSE vs Planet Labs Ref: {ndwi_p:.2f} m") - print(f" RMSE vs USGS Coastlines: {ndwi_u:.2f} m") - print(f" RMSE vs Manual Hi-Res GT: {ndwi_h:.2f} m") - - if dwm_model is not None: - print(" [DeepWaterMap Model]") - print(f" RMSE vs Planet Labs Ref: {dwm_p:.2f} m") - print(f" RMSE vs USGS Coastlines: {dwm_u:.2f} m") - print(f" RMSE vs Manual Hi-Res GT: {dwm_h:.2f} m") - - print("\n Regional RMSE Breakdown vs Planet Labs Ref (U-Net vs NDWI vs DeepWaterMap):") - for r_id in range(1, 6): - r_unet = f"{regional_rmse.get(r_id, float('nan')):.2f} m" if not np.isnan(regional_rmse.get(r_id, float('nan'))) else "N/A" - r_ndwi = f"{ndwi_regional.get(r_id, float('nan')):.2f} m" if not np.isnan(ndwi_regional.get(r_id, float('nan'))) else "N/A" - r_dwm = f"{dwm_regional.get(r_id, float('nan')):.2f} m" if not np.isnan(dwm_regional.get(r_id, float('nan'))) else "N/A" - print(f" - {region_names[r_id]} (R{r_id}): U-Net={r_unet} | NDWI={r_ndwi} | DeepWaterMap={r_dwm}") - - plot_out_path = os.path.join(vis_dir, f"{tile_name}_model_predicted_vs_gt_comparison.png") - plot_predictions_comparison(t_path, pred_lines, ndwi_lines, dwm_lines, planet_union, hires_union, usgs_union, transform, plot_out_path) - - water_land_out_path = os.path.join(vis_dir, f"{tile_name}_model_water_land_prediction.png") - plot_water_land_prediction(t_path, pred_mask_np, pred_lines, water_land_out_path) - -if __name__ == "__main__": - main() diff --git a/training_pipeline/train_and_eval_pipeline.py b/training_pipeline/train_and_eval_pipeline.py deleted file mode 100644 index 7df2ec8..0000000 --- a/training_pipeline/train_and_eval_pipeline.py +++ /dev/null @@ -1,748 +0,0 @@ -import os -import sys -import glob -import re -import numpy as np -import geopandas as gpd -import rasterio as rio -import torch -import torch.nn as nn -import torch.optim as optim -from torch.utils.data import Dataset, DataLoader, random_split -from torchvision import transforms -from PIL import Image -from tqdm import tqdm -import skimage.measure -from shapely.geometry import Point, MultiPoint, LineString, MultiLineString, box -import matplotlib.pyplot as plt - -# Dynamic linker fix: import rasterio and numpy before torch -# Already done at the top. - -# Set path for loading configuration -script_dir = os.path.dirname(os.path.abspath(__file__)) -repo_root = os.path.abspath(os.path.join(script_dir, "..")) -if repo_root not in sys.path: - sys.path.append(repo_root) - -from load_config import load_config, get_augment_tiles_output_folder, get_training_config - -UTM_ZONE_3N = 'EPSG:32603' - -# ---------------------------- -# Dataset Class -# ---------------------------- -class SegmentationDataset(Dataset): - def __init__(self, data_dir, transform=None): - self.data_dir = data_dir - self.transform = transform - self.image_mask_pairs = self._find_image_mask_pairs() - - def _find_image_mask_pairs(self): - pairs = [] - image_files = glob.glob(os.path.join(self.data_dir, "*.tif")) - image_files = [f for f in image_files if "_concatenated_ndwi_mask_" not in os.path.basename(f)] - - for img_path in image_files: - img_name = os.path.basename(img_path) - mask_name = img_name.replace("_clip_", "_concatenated_ndwi_mask_clip_") - mask_path = os.path.join(self.data_dir, mask_name) - - if os.path.exists(mask_path): - pairs.append((img_path, mask_path)) - continue - - base_match = re.match(r'(.+)_\d+-of-\d+(_[^_]+)?\.tif$', img_name) - if base_match: - base_name = base_match.group(1) - mask_name_alt = f"{base_name}_concatenated_ndwi_mask_{img_name.split('_')[-2]}_{img_name.split('_')[-1]}" - mask_path_alt = os.path.join(self.data_dir, mask_name_alt) - if os.path.exists(mask_path_alt): - pairs.append((img_path, mask_path_alt)) - return pairs - - def __len__(self): - return len(self.image_mask_pairs) - - def __getitem__(self, idx): - img_path, mask_path = self.image_mask_pairs[idx] - with rio.open(img_path) as src: - image_data = src.read([3, 2, 1]) - image_data = (np.clip(image_data.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) - image_data = np.transpose(image_data, (1, 2, 0)) - image = Image.fromarray(image_data) - - with rio.open(mask_path) as src: - mask_data = src.read(1) - mask_data = (mask_data > 0).astype(np.uint8) * 255 - mask = Image.fromarray(mask_data, mode="L") - - if self.transform: - image = self.transform(image) - mask = self.transform(mask) - - mask = (mask > 0).float() - return image, mask - -# ---------------------------- -# Model Components -# ---------------------------- -def _init_weights(m): - if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - -class DoubleConv(nn.Module): - """ - Deeper convolution block with 3 convolutional layers, GroupNorm, and residual connection. - """ - def __init__(self, in_channels, out_channels): - super().__init__() - num_groups = min(32, out_channels) - self.conv = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False), - nn.GroupNorm(num_groups, out_channels), - nn.ReLU(inplace=True), - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), - nn.GroupNorm(num_groups, out_channels), - nn.ReLU(inplace=True), - nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), - nn.GroupNorm(num_groups, out_channels), - ) - if in_channels != out_channels: - self.shortcut = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), - nn.GroupNorm(num_groups, out_channels) - ) - else: - self.shortcut = nn.Identity() - self.relu = nn.ReLU(inplace=True) - - def forward(self, x): - return self.relu(self.conv(x) + self.shortcut(x)) - -class UNet(nn.Module): - def __init__(self, n_channels=3, n_classes=1): - super().__init__() - self.down1 = DoubleConv(n_channels, 64) - self.pool1 = nn.MaxPool2d(2) - self.down2 = DoubleConv(64, 128) - self.pool2 = nn.MaxPool2d(2) - self.down3 = DoubleConv(128, 256) - self.pool3 = nn.MaxPool2d(2) - self.down4 = DoubleConv(256, 512) - self.pool4 = nn.MaxPool2d(2) - self.down5 = DoubleConv(512, 1024) - self.pool5 = nn.MaxPool2d(2) - - self.middle = DoubleConv(1024, 2048) - - self.up5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2) - self.conv5 = DoubleConv(2048, 1024) - self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) - self.conv4 = DoubleConv(1024, 512) - self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) - self.conv3 = DoubleConv(512, 256) - self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) - self.conv2 = DoubleConv(256, 128) - self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) - self.conv1 = DoubleConv(128, 64) - - self.final = nn.Conv2d(64, n_classes, kernel_size=1) - self.apply(_init_weights) - - def forward(self, x): - d1 = self.down1(x) - d2 = self.down2(self.pool1(d1)) - d3 = self.down3(self.pool2(d2)) - d4 = self.down4(self.pool3(d3)) - d5 = self.down5(self.pool4(d4)) - mid = self.middle(self.pool5(d5)) - - u5_out = self.up5(mid) - if u5_out.size()[2:] != d5.size()[2:]: - u5_out = nn.functional.interpolate(u5_out, size=d5.size()[2:], mode='bilinear', align_corners=True) - u5 = self.conv5(torch.cat([u5_out, d5], dim=1)) - - u4_out = self.up4(u5) - if u4_out.size()[2:] != d4.size()[2:]: - u4_out = nn.functional.interpolate(u4_out, size=d4.size()[2:], mode='bilinear', align_corners=True) - u4 = self.conv4(torch.cat([u4_out, d4], dim=1)) - - u3_out = self.up3(u4) - if u3_out.size()[2:] != d3.size()[2:]: - u3_out = nn.functional.interpolate(u3_out, size=d3.size()[2:], mode='bilinear', align_corners=True) - u3 = self.conv3(torch.cat([u3_out, d3], dim=1)) - - u2_out = self.up2(u3) - if u2_out.size()[2:] != d2.size()[2:]: - u2_out = nn.functional.interpolate(u2_out, size=d2.size()[2:], mode='bilinear', align_corners=True) - u2 = self.conv2(torch.cat([u2_out, d2], dim=1)) - - u1_out = self.up1(u2) - if u1_out.size()[2:] != d1.size()[2:]: - u1_out = nn.functional.interpolate(u1_out, size=d1.size()[2:], mode='bilinear', align_corners=True) - u1 = self.conv1(torch.cat([u1_out, d1], dim=1)) - - return self.final(u1) # Return raw logits for BCEWithLogitsLoss - -# ---------------------------- -# Attention U-Net -# ---------------------------- -class AttentionGate(nn.Module): - def __init__(self, F_g, F_l, F_int): - super().__init__() - g_groups = min(32, F_int) - x_groups = min(32, F_int) - self.W_g = nn.Sequential( - nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=False), - nn.GroupNorm(g_groups, F_int) - ) - self.W_x = nn.Sequential( - nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=False), - nn.GroupNorm(x_groups, F_int) - ) - self.psi = nn.Sequential( - nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=False), - nn.GroupNorm(1, 1), - nn.Sigmoid() - ) - self.relu = nn.ReLU(inplace=True) - - def forward(self, g, x): - g1 = self.W_g(g) - x1 = self.W_x(x) - if g1.size()[2:] != x1.size()[2:]: - g1 = nn.functional.interpolate(g1, size=x1.size()[2:], mode='bilinear', align_corners=True) - out = self.relu(g1 + x1) - out = self.psi(out) - return x * out - -class AttentionUNet(nn.Module): - def __init__(self, n_channels=3, n_classes=1): - super().__init__() - self.down1 = DoubleConv(n_channels, 64) - self.pool1 = nn.MaxPool2d(2) - self.down2 = DoubleConv(64, 128) - self.pool2 = nn.MaxPool2d(2) - self.down3 = DoubleConv(128, 256) - self.pool3 = nn.MaxPool2d(2) - self.down4 = DoubleConv(256, 512) - self.pool4 = nn.MaxPool2d(2) - self.down5 = DoubleConv(512, 1024) - self.pool5 = nn.MaxPool2d(2) - - self.middle = DoubleConv(1024, 2048) - - self.up5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2) - self.attn5 = AttentionGate(F_g=1024, F_l=1024, F_int=512) - self.conv5 = DoubleConv(2048, 1024) - - self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) - self.attn4 = AttentionGate(F_g=512, F_l=512, F_int=256) - self.conv4 = DoubleConv(1024, 512) - - self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) - self.attn3 = AttentionGate(F_g=256, F_l=256, F_int=128) - self.conv3 = DoubleConv(512, 256) - - self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) - self.attn2 = AttentionGate(F_g=128, F_l=128, F_int=64) - self.conv2 = DoubleConv(256, 128) - - self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) - self.attn1 = AttentionGate(F_g=64, F_l=64, F_int=32) - self.conv1 = DoubleConv(128, 64) - - self.final = nn.Conv2d(64, n_classes, kernel_size=1) - self.apply(_init_weights) - - def forward(self, x): - d1 = self.down1(x) - d2 = self.down2(self.pool1(d1)) - d3 = self.down3(self.pool2(d2)) - d4 = self.down4(self.pool3(d3)) - d5 = self.down5(self.pool4(d4)) - mid = self.middle(self.pool5(d5)) - - up5_out = self.up5(mid) - if up5_out.size()[2:] != d5.size()[2:]: - up5_out = nn.functional.interpolate(up5_out, size=d5.size()[2:], mode='bilinear', align_corners=True) - attn5_out = self.attn5(g=up5_out, x=d5) - u5 = self.conv5(torch.cat([up5_out, attn5_out], dim=1)) - - up4_out = self.up4(u5) - if up4_out.size()[2:] != d4.size()[2:]: - up4_out = nn.functional.interpolate(up4_out, size=d4.size()[2:], mode='bilinear', align_corners=True) - attn4_out = self.attn4(g=up4_out, x=d4) - u4 = self.conv4(torch.cat([up4_out, attn4_out], dim=1)) - - up3_out = self.up3(u4) - if up3_out.size()[2:] != d3.size()[2:]: - up3_out = nn.functional.interpolate(up3_out, size=d3.size()[2:], mode='bilinear', align_corners=True) - attn3_out = self.attn3(g=up3_out, x=d3) - u3 = self.conv3(torch.cat([up3_out, attn3_out], dim=1)) - - up2_out = self.up2(u3) - if up2_out.size()[2:] != d2.size()[2:]: - up2_out = nn.functional.interpolate(up2_out, size=d2.size()[2:], mode='bilinear', align_corners=True) - attn2_out = self.attn2(g=up2_out, x=d2) - u2 = self.conv2(torch.cat([up2_out, attn2_out], dim=1)) - - up1_out = self.up1(u2) - if up1_out.size()[2:] != d1.size()[2:]: - up1_out = nn.functional.interpolate(up1_out, size=d1.size()[2:], mode='bilinear', align_corners=True) - attn1_out = self.attn1(g=up1_out, x=d1) - u1 = self.conv1(torch.cat([up1_out, attn1_out], dim=1)) - - return self.final(u1) # Return raw logits for BCEWithLogitsLoss - -# ---------------------------- -# Loss Functions -# ---------------------------- -class DiceLoss(nn.Module): - def __init__(self, smooth=1e-5): - super().__init__() - self.smooth = smooth - - def forward(self, logits, targets): - probs = torch.sigmoid(logits) - probs = probs.view(-1) - targets = targets.view(-1) - intersection = (probs * targets).sum() - dice = (2. * intersection + self.smooth) / (probs.sum() + targets.sum() + self.smooth) - return 1. - dice - -class CombinedLoss(nn.Module): - def __init__(self, bce_weight=0.5, dice_weight=0.5): - super().__init__() - self.bce = nn.BCEWithLogitsLoss() - self.dice = DiceLoss() - self.bce_weight = bce_weight - self.dice_weight = dice_weight - - def forward(self, logits, targets): - bce_val = self.bce(logits, targets) - dice_val = self.dice(logits, targets) - total_val = self.bce_weight * bce_val + self.dice_weight * dice_val - return total_val, bce_val, dice_val - -# ---------------------------- -# Training Loop Helper -# ---------------------------- -def run_training(model, train_loader, val_loader, epochs, lr, device, save_path): - criterion = CombinedLoss() - optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-2) - scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6) - best_val_loss = float('inf') - - for epoch in range(epochs): - model.train() - train_loss, train_bce, train_dice = 0.0, 0.0, 0.0 - for imgs, masks in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", file=sys.stdout): - imgs, masks = imgs.to(device), masks.to(device) - outputs = model(imgs) - loss, b_loss, d_loss = criterion(outputs, masks) - optimizer.zero_grad() - loss.backward() - optimizer.step() - train_loss += loss.item() - train_bce += b_loss.item() - train_dice += d_loss.item() - - model.eval() - val_loss, val_bce, val_dice = 0.0, 0.0, 0.0 - with torch.no_grad(): - for imgs, masks in val_loader: - imgs, masks = imgs.to(device), masks.to(device) - outputs = model(imgs) - loss, b_loss, d_loss = criterion(outputs, masks) - val_loss += loss.item() - val_bce += b_loss.item() - val_dice += d_loss.item() - - avg_train = train_loss / len(train_loader) - avg_train_bce = train_bce / len(train_loader) - avg_train_dice = train_dice / len(train_loader) - - avg_val = val_loss / len(val_loader) - avg_val_bce = val_bce / len(val_loader) - avg_val_dice = val_dice / len(val_loader) - - current_lr = optimizer.param_groups[0]['lr'] - print(f"Epoch {epoch+1}/{epochs}:") - print(f" Train Loss = {avg_train:.4f} (BCE: {avg_train_bce:.4f}, Dice: {avg_train_dice:.4f})") - print(f" Val Loss = {avg_val:.4f} (BCE: {avg_val_bce:.4f}, Dice: {avg_val_dice:.4f}) | LR = {current_lr:.6f}") - - if avg_val < best_val_loss: - best_val_loss = avg_val - torch.save(model.state_dict(), save_path) - print(f" Saved best model with Val Loss: {avg_val:.4f}") - scheduler.step() - sys.stdout.flush() - - return best_val_loss - -def evaluate_metrics(model, dataloader, device): - model.eval() - tp, fp, fn, tn = 0, 0, 0, 0 - with torch.no_grad(): - for imgs, masks in dataloader: - imgs, masks = imgs.to(device), masks.to(device) - preds = (torch.sigmoid(model(imgs)) > 0.5).float() - tp += ((preds == 1) & (masks == 1)).sum().item() - fp += ((preds == 1) & (masks == 0)).sum().item() - fn += ((preds == 0) & (masks == 1)).sum().item() - tn += ((preds == 0) & (masks == 0)).sum().item() - - total = tp + fp + fn + tn - pixel_acc = (tp + tn) / total if total > 0 else 0 - precision = tp / (tp + fp) if (tp + fp) > 0 else 0 - recall = tp / (tp + fn) if (tp + fn) > 0 else 0 - f1_score = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 - iou = tp / (tp + fp + fn) if (tp + fp + fn) > 0 else 0 - return pixel_acc, precision, recall, f1_score, iou - -# ---------------------------- -# Georeferenced Contours & RMSE -# ---------------------------- -def extract_model_coastline(model, image_path, transform, device): - model.eval() - with rio.open(image_path) as src: - image_data = src.read([3, 2, 1]) - # Scale to 0-255 - image_data = (np.clip(image_data.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) - image_data = np.transpose(image_data, (1, 2, 0)) - h_orig, w_orig = image_data.shape[0], image_data.shape[1] - - image = Image.fromarray(image_data) - transform_resize = transforms.Compose([ - transforms.Resize((256, 256)), - transforms.ToTensor(), - ]) - img_tensor = transform_resize(image).unsqueeze(0).to(device) - - with torch.no_grad(): - output = torch.sigmoid(model(img_tensor)).squeeze().cpu().numpy() - pred_mask = (output > 0.5).astype("uint8") - - # Resize back to original size - pred_mask_resized = Image.fromarray(pred_mask * 255).resize((w_orig, h_orig), Image.NEAREST) - pred_mask_np = np.array(pred_mask_resized) > 0 - - contours = skimage.measure.find_contours(pred_mask_np.astype(np.float32), 0.5) - - lines = [] - for contour in contours: - col = contour[:, 1] - row = contour[:, 0] - xs, ys = rio.transform.xy(transform, row, col) - if len(xs) >= 2: - lines.append(LineString(list(zip(xs, ys)))) - return lines - -def calculate_rmse_on_transects(predicted_lines, trans_deering, ref_distances, usgs_distances): - if not predicted_lines: - return float('nan'), float('nan'), {} - - combined_lines = MultiLineString(predicted_lines) - errors_planet = [] - errors_usgs = [] - regional_errors = {r: [] for r in [1, 2, 3, 4, 5]} - - for idx, row in trans_deering.iterrows(): - oid = int(row['TransOrder']) - t_geom = row.geometry - - # Determine region of transect - region_id = 5 - if oid >= 17443: - region_id = 1 - elif oid >= 17394: - region_id = 2 - elif oid >= 17370: - region_id = 3 - elif oid >= 17337: - region_id = 4 - - pt_int = combined_lines.intersection(t_geom) - dist = None - if not pt_int.is_empty: - if isinstance(pt_int, Point): - dist = t_geom.project(pt_int) - elif isinstance(pt_int, MultiPoint): - dist = min([t_geom.project(pt) for pt in pt_int.geoms]) - elif hasattr(pt_int, 'geoms'): - pts = [pt for pt in pt_int.geoms if isinstance(pt, Point)] - if pts: - dist = min([t_geom.project(pt) for pt in pts]) - - if dist is not None: - if oid in ref_distances: - err_p = dist - ref_distances[oid] - errors_planet.append(err_p) - regional_errors[region_id].append(err_p) - if oid in usgs_distances: - err_u = dist - usgs_distances[oid] - errors_usgs.append(err_u) - - rmse_planet = np.sqrt(np.mean(np.square(errors_planet))) if errors_planet else float('nan') - rmse_usgs = np.sqrt(np.mean(np.square(errors_usgs))) if errors_usgs else float('nan') - - regional_rmse = {} - for r, errs in regional_errors.items(): - regional_rmse[r] = np.sqrt(np.mean(np.square(errs))) if errs else float('nan') - - return rmse_planet, rmse_usgs, regional_rmse - -# ---------------------------- -# Main Execution Block -# ---------------------------- -def main(): - print("==================================================") - print("STARTING PIPELINE: TRAINING & EVALUATION (8 EPOCHS)") - print("==================================================") - - config = load_config() - data_dir = get_augment_tiles_output_folder(config) - device = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu") - print(f"Using device: {device}") - - # 1. Prepare Datasets & Loaders - transform = transforms.Compose([ - transforms.Resize((256, 256)), - transforms.ToTensor(), - ]) - dataset = SegmentationDataset(data_dir, transform=transform) - print(f"Total dataset size: {len(dataset)} samples") - - train_size = int(0.8 * len(dataset)) - val_size = len(dataset) - train_size - generator = torch.Generator().manual_seed(42) - train_set, val_set = random_split(dataset, [train_size, val_size], generator=generator) - - training_config = get_training_config(config) - batch_size = training_config.get('batch_size', 4) - print(f"Loaded batch size: {batch_size}") - - train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False) - val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=False) - - output_dir = os.path.join(repo_root, "output_models") - os.makedirs(output_dir, exist_ok=True) - unet_save_path = os.path.join(output_dir, "best_deep_unet_8epochs.pth") - attn_save_path = os.path.join(output_dir, "best_deep_attn_unet_8epochs.pth") - - # 2. Train Standard U-Net - print("\n----------------------------------") - print("TRAINING STANDARD U-NET (8 EPOCHS)") - print("----------------------------------") - unet_model = UNet(n_channels=3, n_classes=1).to(device) - if os.path.exists(unet_save_path): - print(f"Found existing Standard U-Net weights at {unet_save_path}. Skipping training.") - unet_model.load_state_dict(torch.load(unet_save_path, map_location=device)) - else: - unet_best_val = run_training(unet_model, train_loader, val_loader, epochs=20, lr=1e-3, device=device, save_path=unet_save_path) - - # 3. Train Attention U-Net - print("\n----------------------------------") - print("TRAINING ATTENTION U-NET (8 EPOCHS)") - print("----------------------------------") - attn_model = AttentionUNet(n_channels=3, n_classes=1).to(device) - if os.path.exists(attn_save_path): - print(f"Found existing Attention U-Net weights at {attn_save_path}. Skipping training.") - attn_model.load_state_dict(torch.load(attn_save_path, map_location=device)) - else: - attn_best_val = run_training(attn_model, train_loader, val_loader, epochs=20, lr=1e-3, device=device, save_path=attn_save_path) - - # 4. Load Best Weights for Evaluation - unet_model.load_state_dict(torch.load(unet_save_path, map_location=device)) - attn_model.load_state_dict(torch.load(attn_save_path, map_location=device)) - - # Evaluate Validation Metrics - print("\n----------------------------------") - print("EVALUATING SEGMENTATION METRICS ON VAL SET") - print("----------------------------------") - u_acc, u_prec, u_rec, u_f1, u_iou = evaluate_metrics(unet_model, val_loader, device) - a_acc, a_prec, a_rec, a_f1, a_iou = evaluate_metrics(attn_model, val_loader, device) - - print("\n--- Standard U-Net Validation Metrics ---") - print(f"Pixel Accuracy: {u_acc:.4%}") - print(f"Precision (PPV): {u_prec:.4%}") - print(f"Recall (Sensitivity): {u_rec:.4%}") - print(f"F1-Score (Dice Coeff): {u_f1:.4%}") - print(f"Mean IoU (Jaccard): {u_iou:.4%}") - - print("\n--- Attention U-Net Validation Metrics ---") - print(f"Pixel Accuracy: {a_acc:.4%}") - print(f"Precision (PPV): {a_prec:.4%}") - print(f"Recall (Sensitivity): {a_rec:.4%}") - print(f"F1-Score (Dice Coeff): {a_f1:.4%}") - print(f"Mean IoU (Jaccard): {a_iou:.4%}") - sys.stdout.flush() - - # 5. Georeferenced Testing on 2016 Files & RMSE Calculation - print("\n----------------------------------") - print("EVALUATING HISTORICAL 2016 TEST TILES") - print("----------------------------------") - p_transects = os.path.join(repo_root, "USGS_Coastlines", "WestChukchi_exposed_STepr_rates", "WestChukchi_exposed_STepr_rates.shp") - p_hires = os.path.join(repo_root, "ground_truth", "2016_HiRes_Final_Coastline.shp") - - # Candidate paths for digitized Planet Labs reference coastline - planet_candidates = [ - os.path.join(repo_root, "existing_data", "DigitizedCoastlines", "PlanetCoastline_gt", "09_09_2016", "9_9_16_PlanetCoastline.shp"), - os.path.join(repo_root, "..", "existing_data", "DigitizedCoastlines", "PlanetCoastline_gt", "09_09_2016", "9_9_16_PlanetCoastline.shp") - ] - p_planet = next((p for p in planet_candidates if os.path.exists(p)), planet_candidates[0]) - - trans_deering = gpd.read_file(p_transects).to_crs(UTM_ZONE_3N) - trans_deering = trans_deering[trans_deering['BaselineID'] == 117].sort_values('TransOrder') - hires = gpd.read_file(p_hires).to_crs(UTM_ZONE_3N) - planet = gpd.read_file(p_planet).to_crs(UTM_ZONE_3N) - - hires_union = hires.union_all() if hasattr(hires, 'union_all') else hires.unary_union - planet_union = planet.union_all() if hasattr(planet, 'union_all') else planet.unary_union - - # Precompute reference intersection distances - ref_distances = {} - usgs_distances = {} - for idx, row in trans_deering.iterrows(): - oid = int(row['TransOrder']) - t_geom = row.geometry - p_int = t_geom.intersection(planet_union) - if not p_int.is_empty: - if isinstance(p_int, Point): - ref_distances[oid] = t_geom.project(p_int) - elif isinstance(p_int, MultiPoint): - ref_distances[oid] = min([t_geom.project(pt) for pt in p_int.geoms]) - elif hasattr(p_int, 'geoms'): - pts = [pt for pt in p_int.geoms if isinstance(pt, Point)] - if pts: - ref_distances[oid] = min([t_geom.project(pt) for pt in pts]) - - u_int = t_geom.intersection(hires_union) - if not u_int.is_empty: - if isinstance(u_int, Point): - usgs_distances[oid] = t_geom.project(u_int) - elif isinstance(u_int, MultiPoint): - usgs_distances[oid] = min([t_geom.project(pt) for pt in u_int.geoms]) - elif hasattr(u_int, 'geoms'): - pts = [pt for pt in u_int.geoms if isinstance(pt, Point)] - if pts: - usgs_distances[oid] = min([t_geom.project(pt) for pt in pts]) - - # Test tile paths from test_data_4_6_sept - test_tiles = [ - os.path.join(repo_root, "test_data_4_6_sept", "sept_4", "files", "369619_2016-09-04_RE2_3A_Analytic_SR_clip.tif"), - os.path.join(repo_root, "test_data_4_6_sept", "sept_6", "files", "369619_2016-09-06_RE5_3A_Analytic_SR_clip.tif") - ] - - for t_path in test_tiles: - print(f"\nAnalyzing test tile: {os.path.basename(t_path)}") - with rio.open(t_path) as src: - transform = src.transform - - u_lines = extract_model_coastline(unet_model, t_path, transform, device) - a_lines = extract_model_coastline(attn_model, t_path, transform, device) - - # Calculate RMSE scores - u_rmse_p, u_rmse_u, u_regional = calculate_rmse_on_transects(u_lines, trans_deering, ref_distances, usgs_distances) - a_rmse_p, a_rmse_u, a_regional = calculate_rmse_on_transects(a_lines, trans_deering, ref_distances, usgs_distances) - - print("\n >> U-NET RMSE results:") - print(f" vs Planet Labs Reference: {u_rmse_p:.2f} m") - print(f" vs USGS Ground Truth: {u_rmse_u:.2f} m") - print(" Regional RMSE values:") - print(f" Western Region (R1): {u_regional[1]:.2f} m") - print(f" Northern Region (R2): {u_regional[2]:.2f} m") - print(f" Central Region (R3): {u_regional[3]:.2f} m") - print(f" Town Region (R4): {u_regional[4]:.2f} m") - print(f" East Region (R5): {u_regional[5]:.2f} m") - - print("\n >> ATTENTION U-NET RMSE results:") - print(f" vs Planet Labs Reference: {a_rmse_p:.2f} m") - print(f" vs USGS Ground Truth: {a_rmse_u:.2f} m") - print(" Regional RMSE values:") - print(f" Western Region (R1): {a_regional[1]:.2f} m") - print(f" Northern Region (R2): {a_regional[2]:.2f} m") - print(f" Central Region (R3): {a_regional[3]:.2f} m") - print(f" Town Region (R4): {a_regional[4]:.2f} m") - print(f" East Region (R5): {a_regional[5]:.2f} m") - sys.stdout.flush() - - # Plot predicted vs actual coastlines for the first test tile - if "2016-10-15" in t_path: - plot_predictions_comparison(t_path, u_lines, a_lines, planet_union, hires_union, transform) - -def plot_geometry(ax, geom, color, linewidth, label, linestyle="-"): - if geom.is_empty: - return - if isinstance(geom, LineString): - ax.plot(*geom.xy, color=color, linewidth=linewidth, label=label, linestyle=linestyle) - elif isinstance(geom, MultiLineString): - for line in geom.geoms: - ax.plot(*line.xy, color=color, linewidth=linewidth, label=label, linestyle=linestyle) - elif hasattr(geom, "geoms"): - for sub_geom in geom.geoms: - plot_geometry(ax, sub_geom, color, linewidth, label, linestyle) - -def plot_predictions_comparison(t_path, u_lines, a_lines, planet_union, hires_union, transform): - fig, ax = plt.subplots(figsize=(10, 10)) - ax.set_title("Predicted vs. Actual Coastlines (Tile: 2016-10-15)", fontsize=14, fontweight="bold") - - # Read the RGB image to show as background - with rio.open(t_path) as src: - image_data = src.read([3, 2, 1]) - image_data = (np.clip(image_data.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) - image_data = np.transpose(image_data, (1, 2, 0)) - tile_bounds = src.bounds - extent = [tile_bounds.left, tile_bounds.right, tile_bounds.bottom, tile_bounds.top] - - ax.imshow(image_data, extent=extent, alpha=0.8) - - # Plot Planet reference and USGS Ground truth - tile_box = box(*tile_bounds) - - p_cropped = planet_union.intersection(tile_box) - u_cropped = hires_union.intersection(tile_box) - - # Plot Ground Truth using robust helper - plot_geometry(ax, p_cropped, color="orange", linewidth=2.5, label="Planet Labs Reference") - plot_geometry(ax, u_cropped, color="red", linewidth=2.5, label="USGS Ground Truth") - - # Plot standard U-Net predictions - for idx, line in enumerate(u_lines): - lbl = "Standard U-Net Pred" if idx == 0 else "" - ax.plot(*line.xy, color="cyan", linewidth=1.5, linestyle="--", label=lbl) - - # Plot Attention U-Net predictions - for idx, line in enumerate(a_lines): - lbl = "Attention U-Net Pred" if idx == 0 else "" - ax.plot(*line.xy, color="lime", linewidth=2.0, label=lbl) - - # Simplify legend duplicates - handles, labels = ax.get_legend_handles_labels() - by_label = {} - for h, l in zip(handles, labels): - if l: # skip empty labels - by_label[l] = h - ax.legend(by_label.values(), by_label.keys(), loc="upper right") - - script_dir = os.path.dirname(os.path.abspath(__file__)) - repo_root = os.path.abspath(os.path.join(script_dir, "..")) - output_dir = os.path.join(repo_root, "output_models") - os.makedirs(output_dir, exist_ok=True) - plot_out_path = os.path.join(output_dir, f"{os.path.splitext(os.path.basename(t_path))[0]}_predicted_vs_actual.png") - plt.tight_layout() - plt.savefig(plot_out_path, dpi=300) - plt.close() - print(f"Saved visual predicted vs actual comparison plot to: {plot_out_path}") - -if __name__ == "__main__": - main() diff --git a/training_pipeline/train_deepwatermap.py b/training_pipeline/train_deepwatermap.py new file mode 100644 index 0000000..da61838 --- /dev/null +++ b/training_pipeline/train_deepwatermap.py @@ -0,0 +1,129 @@ +""" +Train 4-Channel DeepWaterMap Model from Scratch on Planet Labs Satellite Tiles + +Uses BCEDiceLoss, spatial data augmentations (flips & rotations), and 50 training epochs. + +Usage: + python train_deepwatermap.py +""" + +import os +import sys +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, random_split +from tqdm import tqdm + +# Add parent directory to sys.path to import load_config +script_dir = os.path.dirname(os.path.abspath(__file__)) +repo_root = os.path.abspath(os.path.join(script_dir, "..")) +if repo_root not in sys.path: + sys.path.append(repo_root) + +from load_config import load_config, get_augment_tiles_output_folder +from deepwatermap_model import DeepWaterMap4Chan, SegmentationDataset4Chan, BCEDiceLoss + +def train_deepwatermap(): + config = load_config() + data_dir = get_augment_tiles_output_folder(config) + output_dir = os.path.join(repo_root, "output_models") + os.makedirs(output_dir, exist_ok=True) + + model_save_path = os.path.join(output_dir, "deepwatermap_planetlabs_best.pth") + checkpoint_path = os.path.join(output_dir, "deepwatermap_planetlabs_checkpoint.pth") + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + epochs = 125 + batch_size = 16 + lr = 1e-4 + + print(f"==================================================") + print(f"TRAINING 4-CHANNEL DEEPWATERMAP MODEL ({epochs} EPOCHS)") + print(f"Device: {device}") + print(f"Data Directory: {data_dir}") + print(f"Epochs: {epochs} | Batch Size: {batch_size} | LR: {lr}") + print(f"Loss Function: BCEDiceLoss (Combined BCE + Dice Loss)") + print(f"Augmentations: Enabled (Flips & Rotations)") + print(f"Model Save Path: {model_save_path}") + print(f"==================================================") + + # Instantiate dataset + full_dataset = SegmentationDataset4Chan(data_dir, image_size=(256, 256), is_train=False) + if len(full_dataset) == 0: + print("Error: No 4-channel image-mask pairs found!") + return + + # 80/20 train/val split + train_size = int(0.8 * len(full_dataset)) + val_size = len(full_dataset) - train_size + train_indices, val_indices = random_split(range(len(full_dataset)), [train_size, val_size]) + + train_set = SegmentationDataset4Chan(data_dir, image_size=(256, 256), is_train=True) + train_set.image_mask_pairs = [full_dataset.image_mask_pairs[i] for i in train_indices.indices] + + val_set = SegmentationDataset4Chan(data_dir, image_size=(256, 256), is_train=False) + val_set.image_mask_pairs = [full_dataset.image_mask_pairs[i] for i in val_indices.indices] + + print(f"Dataset split: {len(train_set)} training samples, {len(val_set)} validation samples") + + train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) + val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True) + + model = DeepWaterMap4Chan(in_channels=4).to(device) + criterion = BCEDiceLoss(bce_weight=0.5) + optimizer = optim.Adam(model.parameters(), lr=lr) + scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=4) + + best_val_loss = float('inf') + + for epoch in range(epochs): + model.train() + running_loss = 0.0 + for imgs, masks in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"): + imgs, masks = imgs.to(device), masks.to(device) + + optimizer.zero_grad() + outputs = model(imgs) + loss = criterion(outputs, masks) + loss.backward() + optimizer.step() + + running_loss += loss.item() + + avg_train_loss = running_loss / len(train_loader) + + # Validation phase + model.eval() + val_loss = 0.0 + with torch.no_grad(): + for imgs, masks in val_loader: + imgs, masks = imgs.to(device), masks.to(device) + outputs = model(imgs) + loss = criterion(outputs, masks) + val_loss += loss.item() + + avg_val_loss = val_loss / len(val_loader) + scheduler.step(avg_val_loss) + + current_lr = optimizer.param_groups[0]['lr'] + print(f"Epoch {epoch+1}/{epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | LR: {current_lr:.6f}") + + if avg_val_loss < best_val_loss: + best_val_loss = avg_val_loss + torch.save(model.state_dict(), model_save_path) + print(f" -> Saved new best 4-channel DeepWaterMap model to {model_save_path}") + + # Save final checkpoint + checkpoint_data = { + 'epoch': epochs, + 'model_state_dict': model.state_dict(), + 'optimizer_state_dict': optimizer.state_dict(), + 'best_val_loss': best_val_loss, + } + torch.save(checkpoint_data, checkpoint_path) + print(f"\nTraining completed! Best Validation Loss: {best_val_loss:.4f}") + print(f"Model saved to: {model_save_path}") + +if __name__ == "__main__": + train_deepwatermap() diff --git a/training_pipeline/train_unet.py b/training_pipeline/train_unet.py index 427c2b3..985bc3e 100644 --- a/training_pipeline/train_unet.py +++ b/training_pipeline/train_unet.py @@ -1,16 +1,20 @@ """ -U-Net Training Script for Coastline Segmentation +U-Net & Attention U-Net Training & Model Definitions for Coastline Segmentation -Trains a U-Net model for coastline segmentation on satellite imagery. -Automatically pairs images with masks and uses configurable parameters. +This module provides dataset utilities, model definitions, checkpoint management, +and training routines for coastline extraction. -Usage: python train_unet.py +Supported Model Architectures: +- ClassicUNet: Standard 4-stage U-Net with BatchNorm and 2-conv blocks. +- UNet / DeepUNet: Deeper 5-stage U-Net with GroupNorm, 3-conv residual blocks. +- AttentionUNet: 5-stage U-Net with Attention Gates for spatial feature filtering. -Configuration: All parameters in config_template.json -Data: Images and masks in results_augment_tiles folder -Naming: Images (*_XX-of-YY_[aug].tif) and Masks (*_concatenated_ndwi_mask_XX-of-YY_[aug].tif) +Usage: + python train_unet.py """ +import rasterio as rio +import numpy as np import os import sys import torch @@ -27,7 +31,11 @@ from datetime import datetime # Add parent directory to path to import load_config -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +script_dir = os.path.dirname(os.path.abspath(__file__)) +repo_root = os.path.abspath(os.path.join(script_dir, "..")) +if repo_root not in sys.path: + sys.path.append(repo_root) + from load_config import load_config, get_augment_tiles_output_folder, get_training_config, get_model_save_path # ---------------------------- @@ -36,50 +44,45 @@ class SegmentationDataset(Dataset): """ Dataset for coastline segmentation training. - - Automatically pairs satellite images with their corresponding masks - based on naming convention. Handles augmented data. - + Automatically pairs GeoTIFF satellite image tiles with corresponding binary NDWI water masks. + Args: - data_dir (str): Directory containing images and masks - transform (callable, optional): Transform to apply to data + data_dir (str): Path to directory containing tiled images and masks. + transform (callable, optional): PyTorch transform to apply to images and masks. """ def __init__(self, data_dir, transform=None): - """Initialize dataset with data directory and optional transforms.""" self.data_dir = data_dir self.transform = transform self.image_mask_pairs = self._find_image_mask_pairs() def _find_image_mask_pairs(self): """ - Find matching image and mask pairs based on naming convention. - - Naming: Images (*_XX-of-YY_[aug].tif) -> Masks (*_concatenated_ndwi_mask_XX-of-YY_[aug].tif) - Returns: List of (image_path, mask_path) tuples + Scans data directory for matching image and mask files. + + Returns: + list of tuple: List of (image_path, mask_path) file pairs. """ pairs = [] - - # Get all image files (tiles without mask in name) image_files = glob.glob(os.path.join(self.data_dir, "*.tif")) image_files = [f for f in image_files if "_concatenated_ndwi_mask_" not in os.path.basename(f)] for img_path in image_files: img_name = os.path.basename(img_path) + mask_name = img_name.replace("_clip_", "_concatenated_ndwi_mask_clip_") + mask_path = os.path.join(self.data_dir, mask_name) - # Extract the base name and augmentation suffix - # Pattern: base_name_XX-of-YY_[augmentation].tif + if os.path.exists(mask_path): + pairs.append((img_path, mask_path)) + continue + base_match = re.match(r'(.+)_\d+-of-\d+(_[^_]+)?\.tif$', img_name) if base_match: base_name = base_match.group(1) - augmentation = base_match.group(2) if base_match.group(2) else "" - - # Construct corresponding mask name - mask_name = f"{base_name}_concatenated_ndwi_mask_{img_name.split('_')[-2]}_{img_name.split('_')[-1]}" - mask_path = os.path.join(self.data_dir, mask_name) - - if os.path.exists(mask_path): - pairs.append((img_path, mask_path)) + mask_name_alt = f"{base_name}_concatenated_ndwi_mask_{img_name.split('_')[-2]}_{img_name.split('_')[-1]}" + mask_path_alt = os.path.join(self.data_dir, mask_name_alt) + if os.path.exists(mask_path_alt): + pairs.append((img_path, mask_path_alt)) else: print(f"Warning: Mask not found for {img_name}") @@ -87,48 +90,69 @@ def _find_image_mask_pairs(self): return pairs def __len__(self): - """Return number of image-mask pairs in dataset.""" + """Returns total number of paired samples in dataset.""" return len(self.image_mask_pairs) def __getitem__(self, idx): """ - Get image-mask pair by index. - - Returns: (image_tensor, mask_tensor) - RGB image and binary mask tensors + Loads and preprocesses image and mask at given index. + + Args: + idx (int): Sample index. + + Returns: + tuple: (transformed_image_tensor, binary_mask_tensor) """ img_path, mask_path = self.image_mask_pairs[idx] - - # Load image and convert to RGB - image = Image.open(img_path).convert("RGB") - - # Load mask and convert to grayscale - mask = Image.open(mask_path).convert("L") + with rio.open(img_path) as src: + image_data = src.read([3, 2, 1]) + image_data = (np.clip(image_data.astype(np.float32) / 10000.0, 0.0, 1.0) * 255.0).astype(np.uint8) + image_data = np.transpose(image_data, (1, 2, 0)) + image = Image.fromarray(image_data) + + with rio.open(mask_path) as src: + mask_data = src.read(1) + mask_data = (mask_data > 0).astype(np.uint8) * 255 + mask = Image.fromarray(mask_data, mode="L") if self.transform: image = self.transform(image) mask = self.transform(mask) - # Convert mask to binary (0 or 1) mask = (mask > 0).float() return image, mask # ---------------------------- -# U-Net Model +# Model Components # ---------------------------- -class DoubleConv(nn.Module): +def _init_weights(m): """ - Double convolution block for U-Net architecture. - - Two consecutive 3x3 convolutions with batch normalization and ReLU. - Used in encoder and decoder paths. - + Initializes module parameters using Kaiming Normal initialization + for Conv layers and constant initialization for Normalization layers. + Args: - in_channels (int): Input channels - out_channels (int): Output channels + m (nn.Module): PyTorch module layer. + """ + if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + +class ClassicDoubleConv(nn.Module): + """ + Classic double convolution block consisting of two 3x3 Conv2d layers, + each followed by BatchNorm2d and ReLU activation. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. """ - def __init__(self, in_channels, out_channels): - """Initialize double convolution block.""" super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), @@ -140,26 +164,104 @@ def __init__(self, in_channels, out_channels): ) def forward(self, x): - """Forward pass through double convolution block.""" return self.double_conv(x) +class ClassicUNet(nn.Module): + """ + Classic 4-stage U-Net architecture for semantic segmentation. + + Args: + n_channels (int): Number of input channels (default: 3 for RGB). + n_classes (int): Number of output classes (default: 1 for binary water segmentation). + """ + def __init__(self, n_channels=3, n_classes=1): + super().__init__() + self.down1 = ClassicDoubleConv(n_channels, 64) + self.pool1 = nn.MaxPool2d(2) + self.down2 = ClassicDoubleConv(64, 128) + self.pool2 = nn.MaxPool2d(2) + self.down3 = ClassicDoubleConv(128, 256) + self.pool3 = nn.MaxPool2d(2) + self.down4 = ClassicDoubleConv(256, 512) + self.pool4 = nn.MaxPool2d(2) + + self.middle = ClassicDoubleConv(512, 1024) + + self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) + self.conv4 = ClassicDoubleConv(1024, 512) + self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) + self.conv3 = ClassicDoubleConv(512, 256) + self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) + self.conv2 = ClassicDoubleConv(256, 128) + self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) + self.conv1 = ClassicDoubleConv(128, 64) + + self.final = nn.Conv2d(64, n_classes, kernel_size=1) + + def forward(self, x): + """ + Forward pass for Classic U-Net. + + Returns: + torch.Tensor: Raw logits tensor of shape (batch_size, n_classes, H, W). + """ + d1 = self.down1(x) + d2 = self.down2(self.pool1(d1)) + d3 = self.down3(self.pool2(d2)) + d4 = self.down4(self.pool3(d3)) + mid = self.middle(self.pool4(d4)) + + u4 = self.conv4(torch.cat([self.up4(mid), d4], dim=1)) + u3 = self.conv3(torch.cat([self.up3(u4), d3], dim=1)) + u2 = self.conv2(torch.cat([self.up2(u3), d2], dim=1)) + u1 = self.conv1(torch.cat([self.up1(u2), d1], dim=1)) + + return self.final(u1) # Return raw logits + +class DoubleConv(nn.Module): + """ + Deeper residual convolution block featuring 3 convolutional layers, + GroupNorm normalization, and residual shortcut connections. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + """ + def __init__(self, in_channels, out_channels): + super().__init__() + num_groups = min(32, out_channels) + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(num_groups, out_channels), + nn.ReLU(inplace=True), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(num_groups, out_channels), + nn.ReLU(inplace=True), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(num_groups, out_channels), + ) + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.GroupNorm(num_groups, out_channels) + ) + else: + self.shortcut = nn.Identity() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.relu(self.conv(x) + self.shortcut(x)) class UNet(nn.Module): """ - U-Net architecture for semantic segmentation. - - Encoder-decoder structure with skip connections for precise segmentation. - Architecture: 4 downsampling blocks -> bottleneck -> 4 upsampling blocks -> final layer. - + Deep 5-stage residual U-Net architecture with GroupNorm and residual blocks. + Args: - n_channels (int): Input channels (default: 3 for RGB) - n_classes (int): Output classes (default: 1 for binary segmentation) + n_channels (int): Input image channels (default: 3). + n_classes (int): Number of output segmentation classes (default: 1). """ - def __init__(self, n_channels=3, n_classes=1): - """Initialize U-Net model with encoder-decoder structure.""" super().__init__() - # Encoder path self.down1 = DoubleConv(n_channels, 64) self.pool1 = nn.MaxPool2d(2) self.down2 = DoubleConv(64, 128) @@ -168,11 +270,13 @@ def __init__(self, n_channels=3, n_classes=1): self.pool3 = nn.MaxPool2d(2) self.down4 = DoubleConv(256, 512) self.pool4 = nn.MaxPool2d(2) + self.down5 = DoubleConv(512, 1024) + self.pool5 = nn.MaxPool2d(2) - # Bottleneck - self.middle = DoubleConv(512, 1024) + self.middle = DoubleConv(1024, 2048) - # Decoder path with skip connections + self.up5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2) + self.conv5 = DoubleConv(2048, 1024) self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) self.conv4 = DoubleConv(1024, 512) self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) @@ -182,26 +286,193 @@ def __init__(self, n_channels=3, n_classes=1): self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) self.conv1 = DoubleConv(128, 64) - # Final classification layer self.final = nn.Conv2d(64, n_classes, kernel_size=1) + self.apply(_init_weights) def forward(self, x): - """Forward pass through U-Net with skip connections.""" - # Encoder path + """ + Forward pass for Deep U-Net. + + Returns: + torch.Tensor: Raw logits tensor of shape (batch_size, n_classes, H, W). + """ d1 = self.down1(x) d2 = self.down2(self.pool1(d1)) d3 = self.down3(self.pool2(d2)) d4 = self.down4(self.pool3(d3)) - mid = self.middle(self.pool4(d4)) + d5 = self.down5(self.pool4(d4)) + mid = self.middle(self.pool5(d5)) - # Decoder path with skip connections - u4 = self.conv4(torch.cat([self.up4(mid), d4], dim=1)) - u3 = self.conv3(torch.cat([self.up3(u4), d3], dim=1)) - u2 = self.conv2(torch.cat([self.up2(u3), d2], dim=1)) - u1 = self.conv1(torch.cat([self.up1(u2), d1], dim=1)) + u5_out = self.up5(mid) + if u5_out.size()[2:] != d5.size()[2:]: + u5_out = nn.functional.interpolate(u5_out, size=d5.size()[2:], mode='bilinear', align_corners=True) + u5 = self.conv5(torch.cat([u5_out, d5], dim=1)) + + u4_out = self.up4(u5) + if u4_out.size()[2:] != d4.size()[2:]: + u4_out = nn.functional.interpolate(u4_out, size=d4.size()[2:], mode='bilinear', align_corners=True) + u4 = self.conv4(torch.cat([u4_out, d4], dim=1)) + + u3_out = self.up3(u4) + if u3_out.size()[2:] != d3.size()[2:]: + u3_out = nn.functional.interpolate(u3_out, size=d3.size()[2:], mode='bilinear', align_corners=True) + u3 = self.conv3(torch.cat([u3_out, d3], dim=1)) - # Final output with sigmoid activation - return torch.sigmoid(self.final(u1)) + u2_out = self.up2(u3) + if u2_out.size()[2:] != d2.size()[2:]: + u2_out = nn.functional.interpolate(u2_out, size=d2.size()[2:], mode='bilinear', align_corners=True) + u2 = self.conv2(torch.cat([u2_out, d2], dim=1)) + + u1_out = self.up1(u2) + if u1_out.size()[2:] != d1.size()[2:]: + u1_out = nn.functional.interpolate(u1_out, size=d1.size()[2:], mode='bilinear', align_corners=True) + u1 = self.conv1(torch.cat([u1_out, d1], dim=1)) + + return self.final(u1) # Return raw logits for BCEWithLogitsLoss + +# Alias DeepUNet +DeepUNet = UNet + +class AttentionGate(nn.Module): + """ + Attention Gate module to filter skip connection feature maps based on + gating signals from deeper network layers. + + Args: + F_g (int): Number of feature maps in gating signal tensor. + F_l (int): Number of feature maps in skip connection tensor. + F_int (int): Intermediate channel reduction dimension. + """ + def __init__(self, F_g, F_l, F_int): + super().__init__() + g_groups = min(32, F_int) + x_groups = min(32, F_int) + self.W_g = nn.Sequential( + nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=False), + nn.GroupNorm(g_groups, F_int) + ) + self.W_x = nn.Sequential( + nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=False), + nn.GroupNorm(x_groups, F_int) + ) + self.psi = nn.Sequential( + nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=False), + nn.GroupNorm(1, 1), + nn.Sigmoid() + ) + self.relu = nn.ReLU(inplace=True) + + def forward(self, g, x): + """ + Forward pass for Attention Gate. + + Args: + g (torch.Tensor): Gating signal tensor from deeper layer. + x (torch.Tensor): Skip connection feature map tensor. + + Returns: + torch.Tensor: Attention-weighted feature map. + """ + g1 = self.W_g(g) + x1 = self.W_x(x) + if g1.size()[2:] != x1.size()[2:]: + g1 = nn.functional.interpolate(g1, size=x1.size()[2:], mode='bilinear', align_corners=True) + out = self.relu(g1 + x1) + out = self.psi(out) + return x * out + +class AttentionUNet(nn.Module): + """ + Deep 5-stage residual Attention U-Net architecture. + Uses Attention Gates on skip connections for enhanced spatial feature selection. + + Args: + n_channels (int): Input image channels (default: 3). + n_classes (int): Output segmentation classes (default: 1). + """ + def __init__(self, n_channels=3, n_classes=1): + super().__init__() + self.down1 = DoubleConv(n_channels, 64) + self.pool1 = nn.MaxPool2d(2) + self.down2 = DoubleConv(64, 128) + self.pool2 = nn.MaxPool2d(2) + self.down3 = DoubleConv(128, 256) + self.pool3 = nn.MaxPool2d(2) + self.down4 = DoubleConv(256, 512) + self.pool4 = nn.MaxPool2d(2) + self.down5 = DoubleConv(512, 1024) + self.pool5 = nn.MaxPool2d(2) + + self.middle = DoubleConv(1024, 2048) + + self.up5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2) + self.attn5 = AttentionGate(F_g=1024, F_l=1024, F_int=512) + self.conv5 = DoubleConv(2048, 1024) + + self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2) + self.attn4 = AttentionGate(F_g=512, F_l=512, F_int=256) + self.conv4 = DoubleConv(1024, 512) + + self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2) + self.attn3 = AttentionGate(F_g=256, F_l=256, F_int=128) + self.conv3 = DoubleConv(512, 256) + + self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2) + self.attn2 = AttentionGate(F_g=128, F_l=128, F_int=64) + self.conv2 = DoubleConv(256, 128) + + self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2) + self.attn1 = AttentionGate(F_g=64, F_l=64, F_int=32) + self.conv1 = DoubleConv(128, 64) + + self.final = nn.Conv2d(64, n_classes, kernel_size=1) + self.apply(_init_weights) + + def forward(self, x): + """ + Forward pass for Attention U-Net. + + Returns: + torch.Tensor: Raw logits tensor of shape (batch_size, n_classes, H, W). + """ + d1 = self.down1(x) + d2 = self.down2(self.pool1(d1)) + d3 = self.down3(self.pool2(d2)) + d4 = self.down4(self.pool3(d3)) + d5 = self.down5(self.pool4(d4)) + mid = self.middle(self.pool5(d5)) + + up5_out = self.up5(mid) + if up5_out.size()[2:] != d5.size()[2:]: + up5_out = nn.functional.interpolate(up5_out, size=d5.size()[2:], mode='bilinear', align_corners=True) + attn5_out = self.attn5(g=up5_out, x=d5) + u5 = self.conv5(torch.cat([up5_out, attn5_out], dim=1)) + + up4_out = self.up4(u5) + if up4_out.size()[2:] != d4.size()[2:]: + up4_out = nn.functional.interpolate(up4_out, size=d4.size()[2:], mode='bilinear', align_corners=True) + attn4_out = self.attn4(g=up4_out, x=d4) + u4 = self.conv4(torch.cat([up4_out, attn4_out], dim=1)) + + up3_out = self.up3(u4) + if up3_out.size()[2:] != d3.size()[2:]: + up3_out = nn.functional.interpolate(up3_out, size=d3.size()[2:], mode='bilinear', align_corners=True) + attn3_out = self.attn3(g=up3_out, x=d3) + u3 = self.conv3(torch.cat([up3_out, attn3_out], dim=1)) + + up2_out = self.up2(u3) + if up2_out.size()[2:] != d2.size()[2:]: + up2_out = nn.functional.interpolate(up2_out, size=d2.size()[2:], mode='bilinear', align_corners=True) + attn2_out = self.attn2(g=up2_out, x=d2) + u2 = self.conv2(torch.cat([up2_out, attn2_out], dim=1)) + + up1_out = self.up1(u2) + if up1_out.size()[2:] != d1.size()[2:]: + up1_out = nn.functional.interpolate(up1_out, size=d1.size()[2:], mode='bilinear', align_corners=True) + attn1_out = self.attn1(g=up1_out, x=d1) + u1 = self.conv1(torch.cat([up1_out, attn1_out], dim=1)) + + return self.final(u1) # Return raw logits for BCEWithLogitsLoss # ---------------------------- # Checkpoint Functions @@ -209,17 +480,17 @@ def forward(self, x): def save_training_checkpoint(model, optimizer, epoch, train_loss, val_loss, best_val_loss, config, checkpoint_path): """ - Save training checkpoint with model state and training progress. - + Saves full training checkpoint dictionary to file. + Args: - model: U-Net model to save - optimizer: Optimizer state - epoch: Current epoch number - train_loss: Current training loss - val_loss: Current validation loss - best_val_loss: Best validation loss so far - config: Training configuration - checkpoint_path: Path to save checkpoint + model (nn.Module): Current model state. + optimizer (torch.optim.Optimizer): Current optimizer state. + epoch (int): Current training epoch index. + train_loss (float): Average training loss for current epoch. + val_loss (float): Average validation loss for current epoch. + best_val_loss (float): Historical best validation loss. + config (dict): Active configuration settings. + checkpoint_path (str): File path to write checkpoint file. """ checkpoint_data = { 'epoch': epoch, @@ -231,32 +502,28 @@ def save_training_checkpoint(model, optimizer, epoch, train_loss, val_loss, 'config': config, 'timestamp': datetime.now().isoformat() } - torch.save(checkpoint_data, checkpoint_path) print(f"Checkpoint saved at epoch {epoch}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}") def load_training_checkpoint(checkpoint_path, model, optimizer=None): """ - Load training checkpoint to resume training. - + Loads saved checkpoint weights and state into model and optimizer. + Args: - checkpoint_path: Path to checkpoint file - model: U-Net model to load state into - optimizer: Optimizer to load state into (optional) - + checkpoint_path (str): Path to checkpoint file. + model (nn.Module): Target model instance. + optimizer (torch.optim.Optimizer, optional): Target optimizer instance. + Returns: - tuple: (epoch, train_loss, val_loss, best_val_loss, config) or None if not found + tuple or None: (epoch, train_loss, val_loss, best_val_loss, config) if loaded, else None. """ if not os.path.exists(checkpoint_path): return None - try: checkpoint = torch.load(checkpoint_path, map_location='cpu') - model.load_state_dict(checkpoint['model_state_dict']) if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer_state_dict']) - print(f"Checkpoint loaded from epoch {checkpoint['epoch']}") return (checkpoint['epoch'], checkpoint['train_loss'], @@ -269,18 +536,17 @@ def load_training_checkpoint(checkpoint_path, model, optimizer=None): def get_checkpoint_path(model_save_path, epoch=None): """ - Generate checkpoint file path. - + Constructs standardized checkpoint file path based on model path and epoch. + Args: - model_save_path: Base model save path - epoch: Epoch number for specific checkpoint (optional) - + model_save_path (str): Base model output save path. + epoch (int, optional): Epoch index. + Returns: - str: Checkpoint file path + str: Absolute or relative checkpoint filepath. """ base_dir = os.path.dirname(model_save_path) base_name = os.path.splitext(os.path.basename(model_save_path))[0] - if epoch is not None: return os.path.join(base_dir, f"{base_name}_checkpoint_epoch_{epoch}.pth") else: @@ -291,19 +557,18 @@ def get_checkpoint_path(model_save_path, epoch=None): # ---------------------------- def train_model(model, train_loader, val_loader, config, model_save_path, resume_from_checkpoint=True): """ - Train U-Net model for coastline segmentation with checkpoint support. - - Implements complete training loop with forward pass, loss computation, - backpropagation, and validation. Uses configurable parameters and progress tracking. - Supports resuming from checkpoints and saving best model. - + Main training loop for U-Net & Attention U-Net models. + + Handles forward pass, BCEWithLogits loss computation, backpropagation, + validation evaluation, checkpointing, and saving the best performing model. + Args: - model (UNet): U-Net model to train - train_loader (DataLoader): Training data loader - val_loader (DataLoader): Validation data loader - config (dict): Configuration with training parameters - model_save_path (str): Path to save trained model - resume_from_checkpoint (bool): Whether to resume from checkpoint if available + model (nn.Module): U-Net model instance. + train_loader (DataLoader): PyTorch training DataLoader. + val_loader (DataLoader): PyTorch validation DataLoader. + config (dict): Configuration parameters dictionary. + model_save_path (str): Filepath to save the best model weights. + resume_from_checkpoint (bool): Whether to resume training from existing checkpoint. """ training_config = get_training_config(config) epochs = training_config.get('epochs', 30) @@ -311,33 +576,29 @@ def train_model(model, train_loader, val_loader, config, model_save_path, resume device = training_config.get('device', 'auto') save_every_n_epochs = training_config.get('save_every_n_epochs', 5) - # Set device if device == 'auto': device = "cuda" if torch.cuda.is_available() else "cpu" - criterion = nn.BCELoss() + criterion = nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=lr) model.to(device) print(f"Training on device: {device}") - # Initialize training variables start_epoch = 0 best_val_loss = float('inf') train_losses = [] val_losses = [] - # Try to load checkpoint if resume is enabled checkpoint_path = get_checkpoint_path(model_save_path) if resume_from_checkpoint: checkpoint_data = load_training_checkpoint(checkpoint_path, model, optimizer) if checkpoint_data is not None: start_epoch, _, _, best_val_loss, _ = checkpoint_data - start_epoch += 1 # Start from next epoch + start_epoch += 1 print(f"Resuming training from epoch {start_epoch}") for epoch in range(start_epoch, epochs): - # Training phase model.train() running_loss = 0.0 for imgs, masks in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"): @@ -351,7 +612,6 @@ def train_model(model, train_loader, val_loader, config, model_save_path, resume optimizer.step() running_loss += loss.item() - # Validation phase val_loss = 0.0 model.eval() with torch.no_grad(): @@ -361,7 +621,6 @@ def train_model(model, train_loader, val_loader, config, model_save_path, resume loss = criterion(outputs, masks) val_loss += loss.item() - # Calculate average losses avg_train_loss = running_loss / len(train_loader) avg_val_loss = val_loss / len(val_loader) train_losses.append(avg_train_loss) @@ -369,67 +628,48 @@ def train_model(model, train_loader, val_loader, config, model_save_path, resume print(f"Epoch {epoch+1}: Train Loss={avg_train_loss:.4f}, Val Loss={avg_val_loss:.4f}") - # Save checkpoint every N epochs if (epoch + 1) % save_every_n_epochs == 0: save_training_checkpoint(model, optimizer, epoch, avg_train_loss, avg_val_loss, best_val_loss, config, checkpoint_path) - # Save best model if avg_val_loss < best_val_loss: best_val_loss = avg_val_loss torch.save(model.state_dict(), model_save_path) print(f"New best model saved! Val Loss: {avg_val_loss:.4f}") - # Final checkpoint save save_training_checkpoint(model, optimizer, epochs-1, avg_train_loss, avg_val_loss, best_val_loss, config, checkpoint_path) print(f"Training completed! Best validation loss: {best_val_loss:.4f}") print(f"Final model saved to {model_save_path}") -# ---------------------------- -# Main -# ---------------------------- if __name__ == "__main__": - """ - Main execution block for U-Net training. - - Handles complete training pipeline: config loading, dataset creation, - train/val split, model training, and saving. Auto-detects GPU/CPU. - """ - # Load configuration config = load_config() training_config = get_training_config(config) - # Get paths from config data_dir = get_augment_tiles_output_folder(config) model_save_path = get_model_save_path(config) - # Get training parameters from config image_size = training_config.get('image_size', [256, 256]) batch_size = training_config.get('batch_size', 8) train_split = training_config.get('train_split', 0.8) + model_type = training_config.get('model_type', 'attention_unet') print(f"Data directory: {data_dir}") print(f"Model save path: {model_save_path}") - print(f"Image size: {image_size}") - print(f"Batch size: {batch_size}") - print(f"Train split: {train_split}") + print(f"Model architecture type: {model_type}") - # Create transforms transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), ]) - # Create dataset dataset = SegmentationDataset(data_dir, transform=transform) if len(dataset) == 0: print("No image-mask pairs found! Please check your data directory and file naming.") - exit(1) + sys.exit(1) - # Split into train and validation sets total_size = len(dataset) train_size = int(train_split * total_size) val_size = total_size - train_size @@ -437,10 +677,14 @@ def train_model(model, train_loader, val_loader, config, model_save_path, resume print(f"Dataset split: {train_size} training, {val_size} validation samples") - # Create data loaders - train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True) - val_loader = DataLoader(val_set, batch_size=batch_size) + train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) + val_loader = DataLoader(val_set, batch_size=batch_size, num_workers=4, pin_memory=True) - # Create and train model - model = UNet(n_channels=3, n_classes=1) + if model_type == 'attention_unet': + model = AttentionUNet(n_channels=3, n_classes=1) + elif model_type == 'classic_unet': + model = ClassicUNet(n_channels=3, n_classes=1) + else: + model = UNet(n_channels=3, n_classes=1) + train_model(model, train_loader, val_loader, config, model_save_path)