diff --git a/.coveragerc b/.coveragerc index 7de6bbb..aafc3c6 100644 --- a/.coveragerc +++ b/.coveragerc @@ -17,4 +17,4 @@ partial_branches = show_missing = True -fail_under = 60 +fail_under = 96 diff --git a/src/aws/osml/image_processing/gdal_tile_factory.py b/src/aws/osml/image_processing/gdal_tile_factory.py index 75b3081..7ae87f3 100644 --- a/src/aws/osml/image_processing/gdal_tile_factory.py +++ b/src/aws/osml/image_processing/gdal_tile_factory.py @@ -246,20 +246,53 @@ def find_appropriate_r_level(src_bbox, tile_width) -> int: map2 = src_y_interpolator(dst_x, dst_y).astype(np.float32) logger.debug( - f"Sanity check remap array sizes. " f"They should match the desired map tile size {tile_size[0]}x{tile_size[1]}" + f"Sanity check remap array sizes. They should match the desired map tile size {tile_size[0]}x{tile_size[1]}" ) logger.debug(f"map1.shape = {map1.shape}") logger.debug(f"map2.shape = {map2.shape}") - dst = cv2.remap(src, map1, map2, cv2.INTER_LINEAR) + # Set the scalar for out of bounds pixels to 0 and bias the image just slightly such that all values + # of 0 in the output are only from out of bounds pixels. Valid range for the scalar is 0-255 [int|tuple]. + scalar = 0 + if src.ndim == 2: # 1-band grayscale + src[src == 0] = 1 + elif src.ndim >= 3 and src.shape[2] == 3: # 3-band + scalar = (0, 0, 0) + src[(src == [0, 0, 0]).all(axis=2)] = [0, 0, 1] + logger.debug(f"scalar = {scalar}") + + # Transform image + dst = cv2.remap(src, map1, map2, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=scalar) + + # Create alpha layer mask + alpha_mask = None + if dst.ndim == 2: # 1-band grayscale + all_channel_pixels_mask = dst != 0 + alpha_mask = np.zeros_like(dst, dtype=np.uint8) + alpha_mask[all_channel_pixels_mask] = 255 + elif dst.ndim >= 3 and dst.shape[2] == 3: # 3-band + all_channel_pixels_mask = np.all(dst != 0, axis=2) + alpha_mask = np.zeros_like(dst[..., 0], dtype=np.uint8) + alpha_mask[all_channel_pixels_mask] = 255 + if alpha_mask is not None: # arrays with zeros/zero size can be falsy so explicitly check None + logger.debug(f"alpha_mask.shape = {alpha_mask.shape}") + elif dst.ndim > 2: + logger.debug(f"alpha_mask = None. Image has {dst.ndim} dimensions and {dst.shape[2]} bands.") + else: + logger.debug(f"alpha_mask = None. Image has {dst.ndim} dimensions.") + output_tile_pixels = self._create_display_image(dst) + if alpha_mask is not None: + # imencode does not support 2-band (grayscale + alpha) so the workaround is to convert to 3-band + if output_tile_pixels.ndim == 2: + output_tile_pixels = np.dstack((output_tile_pixels, output_tile_pixels, output_tile_pixels)) + # add alpha mask + output_tile_pixels = np.dstack((output_tile_pixels, alpha_mask)) + # TODO: Formats other than PNG? is_success, image_bytes = cv2.imencode(".png", output_tile_pixels) - if is_success: - return image_bytes - else: - return None + return image_bytes if is_success else None def _read_from_rlevel_as_array( self, scaled_bbox: Tuple[int, int, int, int], r_level: int, band_numbers: Optional[List[int]] = None diff --git a/test/aws/osml/image_processing/test_gdal_tile_factory.py b/test/aws/osml/image_processing/test_gdal_tile_factory.py index 75056ce..b80f0e9 100644 --- a/test/aws/osml/image_processing/test_gdal_tile_factory.py +++ b/test/aws/osml/image_processing/test_gdal_tile_factory.py @@ -223,6 +223,31 @@ def test_create_map_tiles_for_image(self): assert tile_dataset.RasterYSize == 256 assert tile_dataset.GetDriver().ShortName == GDALImageFormats.PNG + def test_create_map_tiles_for_color_image(self): + tile_set_id = "WebMercatorQuad" + tile_set = MapTileSetFactory.get_for_id(tile_set_id) + full_dataset, sensor_model = load_gdal_dataset("./test/data/small.tif") + tile_factory = GDALTileFactory( + full_dataset, + sensor_model, + GDALImageFormats.PNG, + GDALCompressionOptions.NONE, + output_type=gdalconst.GDT_Byte, + range_adjustment=RangeAdjustmentType.DRA, + ) + tile_matrix = 14 + tile_row = 10830 + tile_col = 8437 + map_tile = tile_set.get_tile(MapTileId(tile_matrix=tile_matrix, tile_row=tile_row, tile_col=tile_col)) + encoded_tile_data = tile_factory.create_orthophoto_tile(geo_bbox=map_tile.bounds, tile_size=map_tile.size) + assert encoded_tile_data is not None + temp_ds_name = "/vsimem/" + token_hex(16) + ".PNG" + gdal.FileFromMemBuffer(temp_ds_name, encoded_tile_data) + tile_dataset = gdal.Open(temp_ds_name) + assert tile_dataset.RasterXSize == 256 + assert tile_dataset.RasterYSize == 256 + assert tile_dataset.GetDriver().ShortName == GDALImageFormats.PNG + if __name__ == "__main__": unittest.main() diff --git a/test/data/small.tif b/test/data/small.tif new file mode 100644 index 0000000..bb0d715 --- /dev/null +++ b/test/data/small.tif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:503276dfbe7ff685f0b7fe212b8abfef1ee7b9dfc2abbcde136aedb16e760740 +size 26111474