# Głęboka segmentacja wodnych dróg śródlądowych Problem segmentacji dróg wodnych dotyczy automatycznego rozróżnienia lustra wody od brzegu oraz innych obiektów. ## Proponowane podejście Zastosowanie architektury U-Net, ze wstępnie wytrenowanym koderem ResNet34. **U-Net** to architektura typu koder-dekoder stworzona w 2015 roku, w celu segmentacji obrazów biomedycznych. Istotą tej architektury są połączenia (*skip-connections*) pomiędzy odpowiednimi warstwami kodera i dekodera. Architektura U-Net znajduje szerokie zastosowanie w zagadnieniach segmentacji, stosowano ją na przykład do detekcji dróg w obrazach lotnicznych, przetwarzania obrazów satelitarnych czy segmentacji otoczenia w pojazdach autonomicznych. W proponowanym modelu U-Net, koder składający się ze zwkyłych warstw splotowych, jest zastąpiony przez bloki sieci ResNet. W sumie koder będzie miał 33 warstwy (z sieci ResNet34 odrzucono ostatnią, w pełni połączoną warstwę). Koder zostanie zainicjalizowany wagami wstępnie wytrenowanymi na zbiorze ImageNet (wykorzystane wagi pochodzą z repozytrium udostępnionego przez PyTorch). ![](https://cdn.hashnode.com/uploads/covers/6accb8d9b7b1651627ad904d/2f51acf0-9ad0-42f1-b886-2a25b3fa3ca7.png align="center") Proponowany model na wejściu przyjmuje kolorowy obraz z kanałami RGB. Aby zapewnić, że wstępnie wytrenowany koder działa poprawnie, zastosowany zostanie preprocessing zgodny ze standardem ImageNet. Obejmuje on skalowanie wartości pikseli do zakresu \[0, 1\] oraz normalizację przy użyciu średniej i odchylenia standardowego dla zbioru ImageNet. Wyjście modelu ma taki sam rozmiar jak obraz wejściowy, ale z jednym kanałem. Na wyjściu modelu zaaplikowana zostanie funkcja aktywacji sigmoid wartości wyjściowe mogą być interpetowane jako prawdopodbieństwo przynależności danego piksela, do klasy, w tym przypadku wody. Funkcja straty modelu oparta jest o współczynnik Sorensena i nosi nazwę *Dice Loss* $$\mathcal{L}_{\text{Dice}} = 1 - \frac{2 \sum_i y_i \hat{y}_i + \epsilon}{\sum_i y_i + \sum_i \hat{y}_i + \epsilon}$$ Strategia uczenia tak zmodyfikowanego modelu u-net moze przebiegać następująco: - W pierwszym etapie wagi enkodera są zamrożone (frozen), uczone są jedynie wagi dekodera. - W następnym etapie kolejne warstwy wag enkodera (zaczynając od góry) mogą być odblokowane i wytrenowane od nowa na zbiorze treningowym, w celu lepszej adaptacji do zadania. ## Zbiór danych Proponowane rozwiązanie zostanie przetestowane na zbiorze danych [LaRS (Lakes, Rivers, Seas)](https://lojzezust.github.io/lars-dataset). Jest to zbiór przeznaczony do panoptycznej segmentacji oraz detekcji obiektów w środowiskach wodnych. Zawiera sceny rejestrowane na jeziorach, rzekach i morzach, dzięki czemu dobrze odzwierciedla zmienne warunki, takie jak różne typy zbiorników, oświetlenie czy pogoda. ## Implementacja Model został zaimplementowany z wykorzystaniem języka Python oraz biblioteki [Segmentation Models PyTroch (SMP)](https://github.com/qubvel-org/segmentation_models.pytorch), która stanowi nakładke na PyTorch, dodającą modele segmentacji. Biblioteka ta jest często wykorzystywana w badaniach czy konkursach dotyczących uczenia maszynowego. Kod zawiera klase generująca obiekt Dataset ```python class WaterBinaryDataset(Dataset): def __init__(self, root, split, image_size, augmentations=None): self.split = split self.image_size = image_size self.augmentations = augmentations self.image_dir = Path(root) / split / "images" self.mask_dir = Path(root) / split / "semantic_masks" self.image_paths = sorted(self.image_dir.glob("*.jpg")) self.has_masks = self.mask_dir.exists() and len(list(self.mask_dir.glob("*.png"))) > 0 if self.has_masks: self.mask_paths = {p.stem: p for p in self.mask_dir.glob("*.png")} self.pairs = [ (img, self.mask_paths[img.stem]) for img in self.image_paths if img.stem in self.mask_paths ] else: # Test images self.pairs = self.image_paths self.resize = T.Resize( size=(image_size, image_size), interpolation=InterpolationMode.BILINEAR, antialias=True, ) print(f"{split}: {len(self.pairs)} samples | masks: {self.has_masks}") def __len__(self): return len(self.pairs) def __getitem__(self, idx): # Test Images if not self.has_masks: img_path = self.pairs[idx] image = read_image(str(img_path)).float() / 255.0 image = tv_tensors.Image(image) image = self.resize(image) if self.augmentations is not None: image = self.augmentations(image) return image.as_subclass(torch.Tensor) # Train and val images img_path, mask_path = self.pairs[idx] image = read_image(str(img_path)).float() / 255.0 mask = read_image(str(mask_path))[0].long() mask[mask == 255] = 0 mask = (mask == 1).float() mask = mask.unsqueeze(0) image = tv_tensors.Image(image) mask = tv_tensors.Mask(mask) image, mask = self.resize(image, mask) if self.augmentations is not None: image, mask = self.augmentations(image, mask) return image.as_subclass(torch.Tensor), mask.as_subclass(torch.Tensor) ``` Definicje modelu w PyTroch, korzystającego dodatkowo z bibilioteki PyLighting ```python class WaterSegmentationModel(pl.LightningModule): def __init__( self, arch, encoder_name, in_channels, out_classes, t_max, lr=2e-4, trainable_backbone_layers=0, **kwargs ): super().__init__() self.save_hyperparameters() self.model = smp.create_model( arch, encoder_name=encoder_name, in_channels=in_channels, classes=out_classes, **kwargs, ) params = smp.encoders.get_preprocessing_params(encoder_name) self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1)) self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1)) self.t_max = t_max self.loss_fn = smp.losses.DiceLoss(smp.losses.BINARY_MODE, from_logits=True) self.lr = lr self.trainable_backbone_layers = trainable_backbone_layers for param in self.model.encoder.parameters(): param.requires_grad = False if trainable_backbone_layers > 0: for layer in list(self.model.encoder.children())[-self.trainable_backbone_layers:]: for param in layer.parameters(): param.requires_grad = True self.training_step_outputs = [] self.validation_step_outputs = [] self.test_step_outputs = [] def forward(self, image): image = (image - self.mean) / self.std return self.model(image) def shared_step(self, batch, stage): image = batch[0] mask = batch[1] logits_mask = self.forward(image) loss = self.loss_fn(logits_mask, mask) prob_mask = logits_mask.sigmoid() pred_mask = (prob_mask > 0.5).float() tp, fp, fn, tn = smp.metrics.get_stats( pred_mask.long(), mask.long(), mode="binary" ) self.log(f"{stage}_loss", loss, on_step=(stage == "train"), on_epoch=True, prog_bar=True) return { "loss": loss, "tp": tp, "fp": fp, "fn": fn, "tn": tn, } def shared_epoch_end(self, outputs, stage): tp = torch.cat([x["tp"] for x in outputs]) fp = torch.cat([x["fp"] for x in outputs]) fn = torch.cat([x["fn"] for x in outputs]) tn = torch.cat([x["tn"] for x in outputs]) # IoU (you already had this) per_image_iou = smp.metrics.iou_score( tp, fp, fn, tn, reduction="micro-imagewise" ) dataset_iou = smp.metrics.iou_score( tp, fp, fn, tn, reduction="micro" ) # Precision, recall, F1 (global / micro) tp_sum = tp.sum().float() fp_sum = fp.sum().float() fn_sum = fn.sum().float() eps = 1e-7 precision = tp_sum / (tp_sum + fp_sum + eps) recall = tp_sum / (tp_sum + fn_sum + eps) f1 = 2 * precision * recall / (precision + recall + eps) self.log_dict( { f"{stage}_per_image_iou": per_image_iou, f"{stage}_dataset_iou": dataset_iou, f"{stage}_precision": precision, f"{stage}_recall": recall, f"{stage}_f1": f1, }, prog_bar=True, ) def training_step(self, batch, batch_idx): out = self.shared_step(batch, "train") self.training_step_outputs.append(out) return out["loss"] def on_train_epoch_end(self): self.shared_epoch_end(self.training_step_outputs, "train") self.training_step_outputs.clear() def validation_step(self, batch, batch_idx): out = self.shared_step(batch, "valid") self.validation_step_outputs.append(out) return out def on_validation_epoch_end(self): self.shared_epoch_end(self.validation_step_outputs, "valid") self.validation_step_outputs.clear() def test_step(self, batch, batch_idx): out = self.shared_step(batch, "test") self.test_step_outputs.append(out) return out def on_test_epoch_end(self): self.shared_epoch_end(self.test_step_outputs, "test") self.test_step_outputs.clear() def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=self.lr) scheduler = lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.t_max, eta_min=1e-5 ) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", "frequency": 1, }, } ```