sat-water 0.1.0__py3-none-any.whl

satwater/inference.py

@@ -0,0 +1,313 @@
+"""
+Created on Wed Nov 18 13:22:57 2023
+
+@author: Busayo Alabi
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+try:
+    import tensorflow as tf
+except Exception as e:
+    raise ImportError(
+        "TensorFlow is required for sat-water inference/training.\n\n"
+        "Install TensorFlow first, then reinstall sat-water.\n"
+        "Recommended:\n"
+        "  Linux/Windows: pip install 'tensorflow'\n"
+        "  Apple Silicon: pip install 'tensorflow-macos' 'tensorflow-metal'\n\n"
+        "If you are using segmentation-models with TF legacy Keras:\n"
+        "  pip install tf-keras segmentation-models\n"
+    ) from e
+
+from satwater.builders import load_pretrained
+from satwater.preprocess import Preprocess
+
+try:
+    from PIL import Image
+except Exception:
+    Image = None
+
+ArrayLikeImage = Union[str, np.ndarray, tf.Tensor, "Image.Image"]
+
+
+@dataclass(frozen=True)
+class SegmentationResult:
+    """
+    Result from a segmentation call.
+    """
+
+    masks: dict[str, np.ndarray]
+    overlays: dict[str, np.ndarray]
+    base_image: np.ndarray
+
+
+class Inference:
+    """
+    Supports both pretrained model keys and local saved model paths:
+    """
+
+    pretrained_keys = {"unet", "resnet34_256", "resnet34_512"}
+
+    def __init__(
+        self, model="unet", *, name=None, repo_id=None, revision="main", save_dir=None
+    ):
+        self.repo_id = repo_id or os.environ.get(
+            "SATWATER_WEIGHTS_REPO", "busayojee/sat-water-weights"
+        )
+        self.revision = revision
+        self.save_dir = Path(save_dir) if save_dir is not None else None
+        self.name = name
+        self.preprocess_func = None
+        self.preprocess_funcs = {}
+
+        if isinstance(model, dict):
+            self.model = {}
+            for key, spec in model.items():
+                m, pp = self._load_one(spec, key_hint=key)
+                self.model[key] = m
+                self.preprocess_funcs[key] = pp
+            return
+
+        if isinstance(model, str):
+            if model in self.pretrained_keys:
+                m, pp = self._load_pretrained(model)
+                self.model = m
+                self.name = model
+                self.preprocess_func = pp
+                return
+
+            if self._looks_like_file(model):
+                self.model = tf.keras.models.load_model(model, compile=False)
+                self.name = name or "local"
+                return
+
+            raise ValueError(
+                "model must be a pretrained key (unet/resnet34_256/resnet34_512), "
+                "a local path to a saved model, or a dict of them."
+            )
+
+        raise TypeError("model must be a str or dict[str, str]")
+
+    @staticmethod
+    def _looks_like_file(s):
+        return ("/" in s) or s.endswith(".keras") or s.endswith(".h5")
+
+    def _load_pretrained(self, key):
+        pm = load_pretrained(
+            model_key=key, repo_id=self.repo_id, revision=self.revision
+        )
+        return pm.model, pm.preprocess_func
+
+    def _load_one(self, spec, key_hint=None):
+        if spec in self.pretrained_keys:
+            return self._load_pretrained(spec)
+        if self._looks_like_file(spec):
+            return tf.keras.models.load_model(spec, compile=False), None
+        if key_hint in self.pretrained_keys:
+            return self._load_pretrained(key_hint)
+        raise ValueError(
+            f"Unknown model spec '{spec}'. Use pretrained key or local model path."
+        )
+
+    @staticmethod
+    def _to_tf_image(x):
+        if isinstance(x, str):
+            raw = tf.io.read_file(x)
+            img = tf.image.decode_image(raw, channels=3, expand_animations=False)
+            img = tf.cast(img, tf.float32) / 255.0
+            return img
+
+        if Image is not None and isinstance(x, Image.Image):
+            arr = np.asarray(x.convert("RGB"), dtype=np.float32) / 255.0
+            return tf.convert_to_tensor(arr, dtype=tf.float32)
+
+        if isinstance(x, np.ndarray):
+            arr = x
+            if arr.ndim == 2:
+                arr = np.stack([arr, arr, arr], axis=-1)
+            if arr.shape[-1] == 4:
+                arr = arr[..., :3]
+            arr = arr.astype(np.float32)
+            if arr.max() > 1.5:
+                arr = arr / 255.0
+            return tf.convert_to_tensor(arr, dtype=tf.float32)
+
+        if isinstance(x, tf.Tensor):
+            t = tf.cast(x, tf.float32)
+            if t.shape.rank == 2:
+                t = tf.stack([t, t, t], axis=-1)
+            if t.shape.rank == 3 and t.shape[-1] == 4:
+                t = t[..., :3]
+            tmax = tf.reduce_max(t)
+            t = tf.cond(tmax > 1.5, lambda: t / 255.0, lambda: t)
+            return t
+
+        raise TypeError(
+            "Unsupported image type. Use path, PIL.Image, numpy array, or tf.Tensor."
+        )
+
+    @staticmethod
+    def _resize(img, size):
+        return tf.image.resize(img, size=size, method="bilinear")
+
+    @staticmethod
+    def _mask_to_uint8(mask):
+        m = np.squeeze(mask)
+        if m.dtype != np.uint8:
+            m = m.astype(np.uint8)
+        return m
+
+    @staticmethod
+    def overlay_mask(image, mask, alpha=0.55, color=(0.0, 0.4, 1.0)):
+        img = image
+        if img.dtype != np.float32:
+            img = img.astype(np.float32) / 255.0
+        img = np.clip(img, 0.0, 1.0)
+
+        m = np.squeeze(mask)
+        water = (m > 0).astype(np.float32)
+        water3 = np.repeat(water[:, :, None], 3, axis=2)
+
+        col = np.array(color, dtype=np.float32)[None, None, :]
+        overlay = img * (1.0 - alpha * water3) + col * (alpha * water3)
+        return (np.clip(overlay, 0.0, 1.0) * 255.0).astype(np.uint8)
+
+    def _maybe_save(self, arr, fname):
+        if self.save_dir is None:
+            return
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+        out = self.save_dir / fname
+        plt.imsave(out.as_posix(), arr)
+
+    def predict(
+        self, image, *, return_overlay=True, save=False, show=False, fname="prediction"
+    ):
+        raw_img = self._to_tf_image(image)
+        model_img = Preprocess.normalization_layer(raw_img)
+        base_img = np.clip(raw_img.numpy(), 0.0, 1.0).astype(np.float32)
+
+        masks = {}
+        overlays = {}
+
+        if isinstance(self.model, dict):
+            for key, model in self.model.items():
+                mask_np = self._predict_for_key(key, model, model_img)
+                masks[key] = mask_np
+                if return_overlay:
+                    ov = self.overlay_mask(
+                        self._image_for_key(key, raw_img).numpy(), mask_np
+                    )
+                    overlays[key] = ov
+                    if save and self.save_dir is not None:
+                        self._maybe_save(ov, f"{fname}_{key}_overlay.png")
+                        self._maybe_save(np.squeeze(mask_np), f"{fname}_{key}_mask.png")
+            if show:
+                self._plot_results(base_img, masks, overlays, title=fname)
+            return SegmentationResult(
+                masks=masks, overlays=overlays, base_image=base_img
+            )
+
+        key = self.name or "model"
+        mask_np = self._predict_for_key(key, self.model, model_img)
+        masks[key] = mask_np
+
+        if return_overlay:
+            ov = self.overlay_mask(self._image_for_key(key, raw_img).numpy(), mask_np)
+            overlays[key] = ov
+            if save and self.save_dir is not None:
+                self._maybe_save(ov, f"{fname}_{key}_overlay.png")
+                self._maybe_save(np.squeeze(mask_np), f"{fname}_{key}_mask.png")
+
+        if show:
+            self._plot_results(base_img, masks, overlays, title=fname)
+
+        return SegmentationResult(masks=masks, overlays=overlays, base_image=base_img)
+
+    def _image_for_key(self, key, img):
+        if key == "unet":
+            return self._resize(img, (128, 128))
+        if key == "resnet34_256":
+            return self._resize(img, (256, 256))
+        if key == "resnet34_512":
+            return self._resize(img, (512, 512))
+        return self._resize(img, (128, 128))
+
+    def _predict_for_key(self, key, model, img):
+        x = self._image_for_key(key, img)
+
+        if isinstance(self.model, dict):
+            pp = self.preprocess_funcs.get(key)
+        else:
+            pp = self.preprocess_func
+
+        if pp is not None:
+            x = pp(x)
+
+        x = (x - tf.reduce_min(x)) / (tf.reduce_max(x) - tf.reduce_min(x) + 1e-8)
+
+        pred = model.predict(x[tf.newaxis, ...], verbose=0)
+        pred = tf.argmax(pred, axis=-1)
+        pred = tf.expand_dims(pred, axis=-1)[0, :, :, :]
+        pred_np = pred.numpy().astype(np.int32)
+
+        if key == "resnet34_512":
+            pred_np = 1 - pred_np
+        return pred_np
+
+    @staticmethod
+    def _plot_results(base_img, masks, overlays, title="result"):
+        keys = list(masks.keys())
+        n = len(keys)
+        cols = 3 if overlays else 2
+        plt.figure(figsize=(5 * cols, 4 * n))
+
+        for i, k in enumerate(keys):
+            r = i * cols
+
+            plt.subplot(n, cols, r + 1)
+            plt.title("Image")
+            plt.imshow(base_img)
+            plt.axis("off")
+
+            plt.subplot(n, cols, r + 2)
+            plt.title(f"Mask: {k}")
+            plt.imshow(np.squeeze(masks[k]))
+            plt.axis("off")
+
+            if overlays:
+                plt.subplot(n, cols, r + 3)
+                plt.title(f"Overlay: {k}")
+                plt.imshow(overlays[k])
+                plt.axis("off")
+
+        plt.suptitle(title)
+        plt.tight_layout()
+        plt.show()
+
+
+def segment_image(
+    image,
+    *,
+    model="unet",
+    repo_id=None,
+    revision="main",
+    return_overlay=True,
+    save=False,
+    save_dir=None,
+    show=False,
+    fname="prediction",
+):
+    infer = Inference(
+        model=model, repo_id=repo_id, revision=revision, save_dir=save_dir
+    )
+    return infer.predict(
+        image, return_overlay=return_overlay, save=save, show=show, fname=fname
+    )

satwater/models.py

@@ -0,0 +1,229 @@
+"""
+Created on Wed Nov 17 11:05:25 2023
+
+@author: Busayo Alabi
+"""
+
+import matplotlib.pyplot as plt
+
+try:
+    import tensorflow as tf
+except Exception as e:
+    raise ImportError(
+        "TensorFlow is required for sat-water inference/training.\n\n"
+        "Install TensorFlow first, then reinstall sat-water.\n"
+        "Recommended:\n"
+        "  Linux/Windows: pip install 'tensorflow'\n"
+        "  Apple Silicon: pip install 'tensorflow-macos' 'tensorflow-metal'\n\n"
+        "If you are using segmentation-models with TF legacy Keras:\n"
+        "  pip install tf-keras segmentation-models\n"
+    ) from e
+
+try:
+    import segmentation_models as sm
+except Exception as e:
+    raise ImportError(
+        "segmentation-models is required for training backbone models.\n"
+        "Install it with:\n"
+        "  pip install segmentation-models tf-keras\n"
+    ) from e
+
+from satwater.preprocess import Preprocess
+
+
+class Unet:
+    loss = sm.losses.categorical_focal_dice_loss
+    metrics = [sm.metrics.iou_score, sm.metrics.f1_score]
+
+    @staticmethod
+    def dice_coef_water(y_true, y_pred, smooth=1e-6):
+        y_true = tf.cast(y_true, tf.float32)
+        y_true = tf.squeeze(y_true, axis=-1)
+        y_pred = tf.cast(y_pred[..., 1], tf.float32)
+
+        y_true_f = tf.reshape(y_true, [tf.shape(y_true)[0], -1])
+        y_pred_f = tf.reshape(y_pred, [tf.shape(y_pred)[0], -1])
+
+        intersection = tf.reduce_sum(y_true_f * y_pred_f, axis=1)
+        denom = tf.reduce_sum(y_true_f + y_pred_f, axis=1)
+
+        dice = (2.0 * intersection + smooth) / (denom + smooth)
+        return tf.reduce_mean(dice)
+
+    @staticmethod
+    def conv_block(input, filters):
+        encoder = tf.keras.layers.Conv2D(filters, (3, 3), padding="same")(input)
+        encoder = tf.keras.layers.BatchNormalization()(encoder)
+        encoder = tf.keras.layers.Activation("relu")(encoder)
+        encoder = tf.keras.layers.Conv2D(filters, (3, 3), padding="same")(encoder)
+        encoder = tf.keras.layers.BatchNormalization()(encoder)
+        encoder = tf.keras.layers.Activation("relu")(encoder)
+        return encoder
+
+    @staticmethod
+    def encoder_block(input, filters):
+        encoder = Unet.conv_block(input, filters)
+        encoder_pool = tf.keras.layers.MaxPooling2D((2, 2), strides=(2, 2))(encoder)
+        return encoder_pool, encoder
+
+    @staticmethod
+    def decoder_block(input, concat, filters):
+        decoder = tf.keras.layers.Conv2DTranspose(
+            filters, (2, 2), strides=(2, 2), padding="same"
+        )(input)
+        decoder = tf.keras.layers.concatenate([concat, decoder], axis=-1)
+        decoder = tf.keras.layers.BatchNormalization()(decoder)
+        decoder = tf.keras.layers.Activation("relu")(decoder)
+        decoder = tf.keras.layers.Conv2D(filters, (3, 3), padding="same")(decoder)
+        decoder = tf.keras.layers.BatchNormalization()(decoder)
+        decoder = tf.keras.layers.Activation("relu")(decoder)
+        decoder = tf.keras.layers.Conv2D(filters, (3, 3), padding="same")(decoder)
+        decoder = tf.keras.layers.BatchNormalization()(decoder)
+        decoder = tf.keras.layers.Activation("relu")(decoder)
+        return decoder
+
+    @staticmethod
+    def models(input_shape, n_classes):
+        inputs = tf.keras.layers.Input(shape=input_shape)
+        encoder_pool0, encoder0 = Unet.encoder_block(inputs, 32)
+        encoder_pool1, encoder1 = Unet.encoder_block(encoder_pool0, 64)
+        encoder_pool2, encoder2 = Unet.encoder_block(encoder_pool1, 128)
+        encoder_pool3, encoder3 = Unet.encoder_block(encoder_pool2, 256)
+        encoder_pool4, encoder4 = Unet.encoder_block(encoder_pool3, 512)
+        center = Unet.conv_block(encoder_pool4, 1024)
+        decoder4 = Unet.decoder_block(center, encoder4, 512)
+        decoder3 = Unet.decoder_block(decoder4, encoder3, 256)
+        decoder2 = Unet.decoder_block(decoder3, encoder2, 128)
+        decoder1 = Unet.decoder_block(decoder2, encoder1, 64)
+        decoder0 = Unet.decoder_block(decoder1, encoder0, 32)
+        outputs = tf.keras.layers.Conv2D(n_classes, (1, 1), activation="softmax")(
+            decoder0
+        )
+        models = tf.keras.models.Model(inputs=[inputs], outputs=[outputs])
+        return models
+
+    @staticmethod
+    def checkpoint():
+        early_stopping = tf.keras.callbacks.EarlyStopping(
+            monitor="val_iou_score", patience=6
+        )
+        reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
+            monitor="val_loss", factor=0.1, patience=4, min_lr=0.00005
+        )
+        return [early_stopping, reduce_lr]
+
+    @staticmethod
+    def train(
+        train_ds,
+        val_ds,
+        epochs=100,
+        shape=(128, 128, 3),
+        n_classes=2,
+        lr=0.0001,
+        loss="sparse_categorical_crossentropy",
+        metrics=None,
+        name="model",
+    ):
+        if metrics is None:
+            metrics = ["accuracy"]
+        model = Unet.models(shape, n_classes)
+        callbacks = Unet.checkpoint()
+        adam = tf.keras.optimizers.legacy.Adam(learning_rate=lr)
+        model.compile(optimizer=adam, loss=loss, metrics=metrics)
+        tf.keras.utils.plot_model(
+            model, to_file=f"segmentation/segment-water/model_{name}.png"
+        )
+        print(model.summary())
+        history = model.fit(
+            train_ds,
+            epochs=100,
+            shuffle=True,
+            verbose=1,
+            callbacks=callbacks,
+            validation_data=val_ds,
+        )
+        model.save(f"{name}.h5")
+        return history
+
+    @staticmethod
+    def plot_history(history, epochs, model=""):
+        IOU = history.history["iou_score"]
+        val_IOU = history.history["val_iou_score"]
+        loss = history.history["loss"]
+        val_loss = history.history["val_loss"]
+        epochs_range = range(epochs)
+        plt.figure(figsize=(16, 8))
+        plt.subplot(1, 2, 1)
+        plt.plot(epochs_range, IOU, label="Training IOU coefficient")
+        plt.plot(epochs_range, val_IOU, label="Validation IOU coefficient")
+        plt.legend(loc="upper right")
+        plt.title("Training and Validation IOU coefficient")
+
+        plt.subplot(1, 2, 2)
+        plt.plot(epochs_range, loss, label="Training Loss")
+        plt.plot(epochs_range, val_loss, label="Validation Loss")
+        plt.legend(loc="upper right")
+        plt.title("Training and Validation Loss")
+        plt.savefig(f"segmentation/segment-water/results/history{model}.png")
+        plt.show()
+
+
+class BackboneModels:
+    parallel_calls = tf.data.AUTOTUNE
+
+    def __init__(self, Backbone, train_ds, val_ds, test_ds, name="model"):
+        assert isinstance(Backbone, str), "Backbone must be string"
+        self.backbone = Backbone
+        self.train_ds = train_ds
+        self.val_ds = val_ds
+        self.test_ds = test_ds
+        self.name = name
+        self.model = None
+
+    @staticmethod
+    def preprocess(images, label, Backbone):
+        preprocessor = sm.get_preprocessing(Backbone)
+        return preprocessor(images), label
+
+    def build_model(self, n_classes, n=1, lr=0.0001):
+        self.train_ds = self.train_ds.map(
+            lambda x, y: BackboneModels.preprocess(x, y, self.backbone),
+            num_parallel_calls=BackboneModels.parallel_calls,
+        )
+        self.val_ds = self.val_ds.map(
+            lambda x, y: BackboneModels.preprocess(x, y, self.backbone),
+            num_parallel_calls=BackboneModels.parallel_calls,
+        )
+        self.test_ds = self.test_ds.map(
+            lambda x, y: BackboneModels.preprocess(x, y, self.backbone),
+            num_parallel_calls=BackboneModels.parallel_calls,
+        )
+        model = sm.Unet(
+            self.backbone,
+            classes=n_classes,
+            activation="softmax",
+            encoder_weights="imagenet",
+        )
+        adam = tf.keras.optimizers.legacy.Adam(learning_rate=lr)
+        model.compile(optimizer=adam, loss=Unet.loss, metrics=Unet.metrics)
+        tf.keras.utils.plot_model(
+            model, to_file=f"segmentation/segment-water/model_{self.name}.png"
+        )
+        self.model = model
+        Preprocess.plot_image(self.train_ds, n)
+        return self.model.summary()
+
+    def train(self):
+        if self.model:
+            history1 = self.model.fit(
+                self.train_ds,
+                epochs=100,
+                verbose=1,
+                callbacks=Unet.checkpoint(),
+                shuffle=True,
+                validation_data=self.val_ds,
+            )
+            self.model.save(f"sat_water_{self.name}.h5")
+            return history1
+        else:
+            return "A model has to be built with the build_model method first"

satwater/preprocess.py

@@ -0,0 +1,161 @@
+"""
+Created on Wed Nov 17 10:48:32 2023
+
+@author: Busayo Alabi
+"""
+
+import random
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+try:
+    import tensorflow as tf
+except Exception as e:
+    raise ImportError(
+        "TensorFlow is required for sat-water inference/training.\n\n"
+        "Install TensorFlow first, then reinstall sat-water.\n"
+        "Recommended:\n"
+        "  Linux/Windows: pip install 'tensorflow'\n"
+        "  Apple Silicon: pip install 'tensorflow-macos' 'tensorflow-metal'\n\n"
+        "If you are using segmentation-models with TF legacy Keras:\n"
+        "  pip install tf-keras segmentation-models\n"
+    ) from e
+
+from skimage import color
+
+
+class Preprocess:
+    normalization_layer = tf.keras.layers.Rescaling(1.0 / 255)
+
+    @staticmethod
+    def load(mask_path, mask_folder, image_folder, channels):
+        mask = tf.io.read_file(mask_path)
+        mask = tf.dtypes.cast(tf.image.decode_jpeg(mask, channels=1), tf.float32)
+        img_path = tf.strings.regex_replace(mask_path, mask_folder, image_folder)
+        image = tf.io.read_file(img_path)
+        image = tf.image.decode_jpeg(image, channels=channels)
+        return {"image": Preprocess.normalization_layer(image), "mask": mask}
+
+    @staticmethod
+    def make_single_class(sample):
+        mask = sample["mask"]
+        sample["mask"] = tf.where(mask >= 0.5, 1, 0)
+        return sample
+
+    @staticmethod
+    def resize(data, size):
+        image = tf.image.resize(data["image"], size[:2])
+        mask = tf.image.resize(data["mask"], size[:2])
+        return image, mask
+
+    def resize_test(data, size):
+        image = tf.image.resize(data, size)
+        return image
+
+    @staticmethod
+    def data_load(
+        dataset, mask_folder, image_folder, split, shape, batch_size, channels=0
+    ):
+        dataset = [dataset]
+        ds = []
+        for f in dataset:
+            mask_path = list(Path(f + mask_folder).glob("**/*"))
+            for i in mask_path:
+                ds.append(str(i))
+        size = len(ds)
+        random.shuffle(ds)
+        ds = tf.data.Dataset.from_tensor_slices(ds)
+        val_size, test_size = split
+        train_size = int((1 - (val_size + test_size)) * size)
+        val_size = int(val_size * size)
+        test_size = int(test_size * size)
+        train_ds = ds.take(train_size)
+        rest_ds = ds.skip(train_size)
+        val_ds = rest_ds.take(val_size)
+        test_ds = rest_ds.skip(val_size)
+
+        AUTOTUNE = tf.data.experimental.AUTOTUNE
+
+        print(f"The Train Dataset contains {len(train_ds)} images.")
+        print(f"The Validation Dataset contains {len(val_ds)} images.")
+        print(f"The Test Dataset contains {len(test_ds)} images.")
+
+        train_ds = train_ds.map(
+            lambda x: Preprocess.load(x, mask_folder, image_folder, channels)
+        )
+        train_ds = train_ds.map(Preprocess.make_single_class)
+        train_ds = train_ds.map(
+            lambda x: Preprocess.resize(x, shape), num_parallel_calls=AUTOTUNE
+        )
+        train_ds = train_ds.batch(batch_size)
+        train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)
+
+        val_ds = val_ds.map(
+            lambda x: Preprocess.load(x, mask_folder, image_folder, channels)
+        )
+        val_ds = val_ds.map(Preprocess.make_single_class)
+        val_ds = val_ds.map(
+            lambda x: Preprocess.resize(x, shape), num_parallel_calls=AUTOTUNE
+        )
+        val_ds = val_ds.batch(batch_size)
+        val_ds = val_ds.prefetch(buffer_size=AUTOTUNE)
+
+        test_ds = test_ds.map(
+            lambda x: Preprocess.load(x, mask_folder, image_folder, channels)
+        )
+        test_ds = test_ds.map(Preprocess.make_single_class)
+        test_ds = test_ds.map(
+            lambda x: Preprocess.resize(x, shape), num_parallel_calls=AUTOTUNE
+        )
+        test_ds = test_ds.batch(batch_size)
+        test_ds = test_ds.prefetch(buffer_size=AUTOTUNE)
+
+        return (train_ds, val_ds, test_ds)
+
+    @staticmethod
+    def label_overlay(img, mask, fname=None):
+        plt.figure()
+
+        mask = np.squeeze(mask, axis=-1)
+
+        overlay = color.label2rgb(
+            mask,
+            img,
+            colors=[(0, 255, 0)],
+            alpha=0.05,
+            bg_label=0,
+            bg_color=None,
+            image_alpha=1,
+            saturation=1,
+            kind="overlay",
+        )
+        plt.figure(figsize=(6, 6))
+
+        plt.imshow(overlay)
+        plt.axis("off")
+
+        if fname:
+            plt.savefig(fname)
+            plt.clf()
+            plt.close()
+
+        plt.show()
+
+    @staticmethod
+    def plot_image(train_ds, n=1, fname=None):
+        for image, mask in (
+            train_ds.unbatch().shuffle(buffer_size=len(train_ds)).take(1)
+        ):
+            image = (image - tf.reduce_min(image)) / (
+                tf.reduce_max(image) - tf.reduce_min(image)
+            )
+
+            plt.figure(figsize=(12, 6))
+            plt.subplot(121)
+            plt.imshow(image)
+            plt.subplot(122)
+            plt.imshow(mask)
+            plt.show()
+            Preprocess.label_overlay(image, mask, fname)