Untitled

@celery_app.task
def apply_mask_on_bands(
    fid: str,
    s3_metadata: dict,
    imagery_metadata: dict,
    ryp_dispatch: bool = False,
) -> None:
    """Applies mask on individual bands, invalidates bad pixels."""
    logger.info("Apply mask on bands task.")
    raster_path = s3_metadata["raster_local_path"]
    overlay_path = s3_metadata["mask_local_path"]

    logger.info("Reading raster masks.")
    with rasterio.open(overlay_path) as src:
        snow_mask = src.read(2).astype(bool)
        shadow_mask = src.read(3).astype(bool)
        cloud_mask = src.read(6).astype(bool)

    mask = snow_mask + shadow_mask + cloud_mask

    with rasterio.open(raster_path, "r+") as src:
        out_meta = src.meta
        profile = src.profile
        src.nodata = np.nan
        img_data = src.read(
            [
                7,
                6,
                5,
            ],
            masked=True,
        )  # apply RGB ordering

    # apply mask
    img_data.mask = mask
    img_data.meta = out_meta.copy()
    img_data = img_data.filled(
        fill_value=np.nan,
    )

    kwargs = src.meta
    kwargs.update(dtype=rasterio.float32, count=3)

    # Write img_data to a new raster file
    with rasterio.open(raster_path, "w", **kwargs) as dst:
        dst.write(img_data.astype(rasterio.float32))

    ...
    ...