Adding bands -take 2

So, folks I need help here. I want to verify if I’m doing this thing right.

We have four bands that we are given , I want to do band-difference and get NDVI and use that as an additional band for training. So, here is where I’m at, and not sure if I have it working.

I’m using - How to Use Deep Learning, PyTorch Lightning, and the Planetary Computer to Predict Cloud Cover in Satellite Imagery - DrivenData Labs code to test it.

In the ClodDataset class, I modified __getitem(…)


    def __getitem__(self, idx: int):
        # Loads an n-channel image from a chip-level dataframe
        img = self.data.loc[idx]
        band_arrs = []
        for band in self.bands:
            with rasterio.open(img[f"{band}_path"]) as b:
                band_arr = b.read(1).astype("float32")
            band_arrs.append(band_arr)
        #--- add Allow division by zero
        np.seterr(divide='ignore', invalid='ignore')

        # --Calculate NDVI
        ndvi = (band_arrs[3].astype("float32") - band_arrs[2].astype("float32")) / (band_arrs[3].astype("float32") + band_arrs[2].astype("float32")).astype("float32") 
        band_arrs.append(ndvi)

        x_arr = np.stack(band_arrs, axis=-1)

Then, in the CloudModel, I changed the number of bands


def _prepare_model(self):
        # Instantiate U-Net model
        unet_model = smp.UnetPlusPlus(
            encoder_name=self.backbone,
            encoder_weights=self.weights,
            in_channels=5,
            classes=2,
            decoder_use_batchnorm=True
        )
        if self.gpu:
            unet_model.cuda()

        return unet_model

Am I on the right track ? Is it the right way to do add bands from the existing 4 bands ?