geowatch.tasks.fusion.coco_stitcher module

Defines the CocoStitchingManager, which stitches predictions from subregions of an image or video (or more generally - data cube) back into rasters corresponding to the original data. This requires that the user use a sliding window (e.g. perhaps defined by kwarray.SlidingWindow) to iterate over space/time, produce predictions, and know the coordinates those predictions should be stitched back together at.

The following attempts to provide a minimal example with a visualization.

CommandLine

xdoctest -m geowatch.tasks.fusion.coco_stitcher __doc__:0

Example

>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> from geowatch.tasks.fusion.coco_stitcher import demo_coco_stitching_manager
>>> # See the contents of the function for details, might port it to a full
>>> # doctest later.
>>> demo_coco_stitching_manager()
geowatch.tasks.fusion.coco_stitcher.demo_coco_stitching_manager()[source]
class geowatch.tasks.fusion.coco_stitcher.CocoStitchingManager(result_dataset, short_code=None, chan_code=None, stiching_space='video', device='numpy', memmap=None, thresh=0.5, write_probs=True, write_preds=False, num_bands='auto', prob_compress='DEFLATE', prob_blocksize=128, prob_format='cog', polygon_categories=None, expected_minmax=None, quantize=True, writer_queue=None, write_prediction_attrs=True, assets_dname='_assets', dtype='float32')[source]

Bases: object

Manage stitching for multiple images / videos in a CocoDataset.

This is done in a memory-efficient way where after all sub-regions in an image or video have been completed, it is finalized, written to the kwcoco manifest / disk, and the memory used for stitching is freed.

Parameters:

  • result_dataset (CocoDataset) – The CocoDataset that is being predicted on. This will be modified when an image prediction is finalized.

  • short_code (str) – short identifier used for directory names. TODO: rename to prefix? OR or something more indicative that this is a directory name?

  • chan_code (str) – If saving the stitched features, this is the channel code to use.

  • stiching_space (str) – Indicates if the results are given in image or video space (up to a scale factor).

  • device (‘numpy’ | torch.device) – Device to stitch on.

  • memmap (bool | PathLike) – if truthy, the stitcher will use a memory map. If this pathlike, then we use this as the directory for the memmap. If True, a temp directory is used.

  • thresh (float) – if making hard decisions, determines the threshold for converting a soft mask into a hard mask, which can be converted into a polygon.

  • prob_compress (str) – Compression algorithm to use when writing probabilities to disk. Can be any GDAL compression code, e.g LZW, DEFLATE, RAW, etc.

  • prob_blocksize (int) – tiled blocksize for output predictions. Defaults to 128.

  • prob_format (str) – the format of the output images. (png, tif, cog).

  • polygon_categories (List[str] | None) – These are the list of channels that should be transformed into polygons. If not set, all are used.

  • quantize (bool) – if True quantize heatmaps before writing them to disk

  • expected_minmax (Tuple[float, float]) – The expected minimum and maximum values allowed in the output to be stitched – i.e. (0, 1) for probabilities. If unspecified this is infered per image.

  • writer_queue (None | BlockingJobQueue) – if specified, uses this shared writer queue, otherwise creates its own.

  • write_prediction_attrs (bool) – set to True if you are adding predictions to the kwcoco file, otherwise set to False to remove unnecessary attributes.

  • dtype (str) – the dtype to stitch over. Defaults to ‘float32’

  • assets_dname (str) – The name of the top-level directory to write new assets. Defaults to _assets

Todo

  • [ ] Handle the case where the input space is related to the output

    space by an affine transform.

  • [X] Handle stitching in image space

  • [X] Handle the case where we are only stitching over images

  • [ ] Handle the case where iteration is non-contiguous, i.e. define

    a robust criterion to determine when an image is “done” being stitched.

  • [ ] Perhaps separate the “soft-probability” prediction stitcher

    from (a) the code that converts soft-to-hard predictions (b) the code that adds hard predictions to the kwcoco file and (c) the code that adds soft predictions to the kwcoco file?

  • [ ] TODO: remove polygon “predictions” from this completely.

Example

>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> import geowatch
>>> dset = geowatch.coerce_kwcoco('geowatch-msi', geodata=True, dates=True,
>>>                            multispectral=True)
>>> result_dataset = dset.copy()
>>> self = CocoStitchingManager(
>>>     result_dataset=result_dataset,
>>>     short_code='demofeat',
>>>     chan_code='df1|df2',
>>>     prob_format='png',
>>>     stiching_space='video')
>>> coco_img = result_dataset.images().coco_images[0]
>>> # Compute a feature in 0.5 video space for a subset of an image
>>> gid = coco_img.img['id']
>>> hidden = coco_img.imdelay(space='video').finalize().mean(axis=2)
>>> my_feature = kwimage.imresize(hidden, scale=0.5)
>>> asset_dsize = my_feature.shape[0:2][::-1]
>>> space_slice = None
>>> self.accumulate_image(gid, space_slice, my_feature,
>>>                       asset_dsize=asset_dsize,
>>>                       scale_asset_from_stitchspace=0.5)
>>> self.finalize_image(gid)
>>> # The new auxiliary image is now in our result dataset
>>> result_img = result_dataset.coco_image(gid)
>>> print(ub.urepr(result_img.img, nl=-1))
>>> assert 'df1' in result_img.channels
>>> im1 = result_img.imdelay('df1', space='video')
>>> im2 = result_img.imdelay(channels='df1', space='asset')
>>> assert im1.shape[0] == hidden.shape[0]
>>> assert im2.shape[0] == my_feature.shape[0]

Example

>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> import geowatch
>>> dset = geowatch.coerce_kwcoco('geowatch-msi', geodata=True, dates=True,
>>>                            multispectral=True)
>>> result_dataset = dset.copy()
>>> self = CocoStitchingManager(
>>>     result_dataset=result_dataset,
>>>     short_code='demofeat',
>>>     chan_code='df1|df2',
>>>     stiching_space='image')
>>> coco_img = result_dataset.images().coco_images[0]
>>> # Compute a feature in 0.5 image space for a subset of an image
>>> gid = coco_img.img['id']
>>> hidden = coco_img.imdelay(space='image').finalize().mean(axis=2)
>>> my_feature = kwimage.imresize(hidden, scale=0.5)
>>> asset_dsize = my_feature.shape[0:2][::-1]
>>> space_slice = None
>>> self.accumulate_image(gid, space_slice, my_feature,
>>>                       asset_dsize=asset_dsize,
>>>                       scale_asset_from_stitchspace=0.5)
>>> self.finalize_image(gid)
>>> # The new auxiliary image is now in our result dataset
>>> result_img = result_dataset.coco_image(gid)
>>> print(ub.urepr(result_img.img, nl=-1))
>>> assert 'df1' in result_img.channels
>>> im1 = result_img.imdelay('df1', space='image')
>>> im2 = result_img.imdelay(channels='df1', space='asset')
>>> assert im1.shape[0] == 600
>>> assert im2.shape[0] == 300
accumulate_image(gid, space_slice, data, asset_dsize=None, scale_asset_from_stitchspace=None, is_ready='auto', weights=None, downweight_edges=False, **kwargs)[source]

Stitches a result into the appropriate image stitcher.

Parameters:

  • gid (int) – the image id to stitch into

  • space_slice (Tuple[slice, slice] | None) – the slice (in “output-space”) the data corresponds to. if None, assumes this is for the entire image.

  • data (ndarray | Tensor) – the feature or probability data

  • asset_dsize (Tuple) – the w/h of outputspace (i.e. the asset we will write)

  • scale_asset_from_stitchspace (float | None) – the scale to the outspace from from the stitch-space (i.e. image/video) space.

  • is_ready (bool) – todo, fix this to work better

  • ** kwargs – dsize, scale deprecated

Note

Output space is asset space for the new asset we are building. The actual stitcher holds data in outspace / assetspace. May want to adjust termonology here.

managed_image_ids()[source]

Return all image ids that are being managed and may be completed or in the process of stitching.

Returns:

image ids

Return type:

List[int]

ready_image_ids()[source]

Returns all image-ids that are known to be ready to finalize.

Returns:

image ids

Return type:

List[int]

submit_finalize_image(gid)[source]

Like finalize image, but submits the job to the manager’s writer queue, which could be asynchronous.

flush_images()[source]

Allow the writer queue to finish finalizing any incomplete images before allowing the process to procede.

property seen_image_ids
finalize_image(gid)[source]

Finalizes the stitcher for this image, deletes it, and adds its hard and/or soft predictions to the CocoDataset.

Parameters:

gid (int) – the image-id to finalize

geowatch.tasks.fusion.coco_stitcher.quantize_image(imdata, old_min=None, old_max=None, quantize_dtype=<class 'numpy.int16'>)[source]

Quantize a float image into an integer representation with the ability to approximately invert back.

Todo

  • [ ] How does this live relative to dequantize in delayed image?

It seems they should be tied somehow.

Parameters:

  • imdata (ndarray) – image data to quantize

  • old_min (float | None) – a stanard floor for minimum values to make quantization consistent across images. If unspecified chooses the minimum value in the data.

  • old_max (float | None) – a stanard ceiling for maximum values to make quantization consistent across images. If unspecified chooses the maximum value in the data.

  • quantize_dtype (dtype) – which type of integer to quantize as

Returns:

Tuple[ndarray, Dict] - new data with encoding information

Note

Setting old_min / old_max indicates the possible extend of the input data (and it will be clipped to it). It does not mean that the input data has to have those min and max values, but it should be between them.

Example

>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> from delayed_image.helpers import dequantize
>>> # Test error when input is not nicely between 0 and 1
>>> imdata = (np.random.randn(32, 32, 3) - 1.) * 2.5
>>> quant1, quantization1 = quantize_image(imdata)
>>> recon1 = dequantize(quant1, quantization1)
>>> error1 = np.abs((recon1 - imdata)).sum()
>>> print('error1 = {!r}'.format(error1))
>>> #
>>> for i in range(1, 20):
>>>     print('i = {!r}'.format(i))
>>>     quant2, quantization2 = quantize_image(imdata, old_min=-i, old_max=i)
>>>     recon2 = dequantize(quant2, quantization2)
>>>     error2 = np.abs((recon2 - imdata)).sum()
>>>     print('error2 = {!r}'.format(error2))

Example

>>> # Test dequantize with uint8
>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> from delayed_image.helpers import dequantize
>>> imdata = np.random.randn(32, 32, 3)
>>> quant1, quantization1 = quantize_image(imdata, quantize_dtype=np.uint8)
>>> recon1 = dequantize(quant1, quantization1)
>>> error1 = np.abs((recon1 - imdata)).sum()
>>> print('error1 = {!r}'.format(error1))

Example

>>> # Test quantization with different signed / unsigned combos
>>> from geowatch.tasks.fusion.coco_stitcher import *  # NOQA
>>> print(quantize_image(None, 0, 1, np.int16))
>>> print(quantize_image(None, 0, 1, np.int8))
>>> print(quantize_image(None, 0, 1, np.uint8))
>>> print(quantize_image(None, 0, 1, np.uint16))
geowatch.tasks.fusion.coco_stitcher.fix_slice(sl)[source]