Reconstruct Regression Image from Patches

[jejjohnson]

Hi, I wanted to know if you all have a built-in method/module for reconstructing an image from patches using real values? I'm thinking in terms of an autoencoder.

# initialize dataset
ds: RegressionDataset = ...

# make predictions
predictions: Float[Array, "B C H W"] = ...

# reconstruct from predictions
ds_labels: Dataset = RegressionLabels.from_predictions(ds.windows, predictions)

I've seen this for the Labels object, in particular the SemanticSegmentation, ObjectDetection and ChipClassification. But I'm thinking even more simply, regression. I feel like there is already something within the package but I'm just missing the keyword from the documentation. Or perhaps there is a more trivial way to do this in the case of regression?

[AdeelH]

Unfortunately, we currently do not have a built-in Labels or LabelStore subclass for regression. Support for regression is limited to the Regression*Dataset classes, the RegressionLearner, and the RegressionVisualizer.

If we were to implement it, I imagine it would look pretty similar to the SemanticSegmentationSmoothLabels class. In fact, it might even be possible to use it directly for regression without any changes.

I think you might be able to get away with the following. Let me know if it works.

ds_labels: Dataset = SemanticSegmentationSmoothLabels.from_predictions(ds.windows, predictions)
ds_labels.save(
    uri='/path/to/pred_dir/',
    crs_transformer=ds.scene.raster_source.crs_transformer,
    class_config=class_config,
    smooth_output=True,
    discrete_output=False,
)

[jejjohnson]

So it does work out (pretty much out of the box). I'll include the pseudo-code below

# get dataset
ds: RegressionSlidingWindowGeoDataset = ...

# create a dataloader
dl = DataLoader(ds, batch_size=32, shuffle=False)

# create dataloader for the predictions
predictions = []

for ibatch in dl:
    # select first element (no labels)
    ibatch = ibatch[0]
   # insert code to apply predictions
    predictions.append(ibatch)

# concat results and change to numpy
predictions: np.ndarray["B C H W"] = torch.cat(predictions).detach().numpy()

# initialize smooth labels dataset
ds_labels: Dataset = SemanticSegmentationSmoothLabels.from_predictions(
   windows: List[Box]=ds.windows, 
   predictions: List[Array]=predictions,
   extent: Box=ds.scene.bbox,
   num_classes=1
)

# extract image directly
img_pred: Array = ds_labels.pixel_scores

# create fake class config
class_config = ClassConfig(names=["blah"])
# save to tif image
ds_labels.save(
    "./temp", 
    class_config: ClassConfig =class_config,
    crs_transformer: CRSTransformer=ds.scene.raster_source.crs_transformer,
    smooth_output: bool=True,
    discrete_output: bool=False,
)

---

Tips n Tricks

Sliding Window Dataset

We need to use the sliding window dataset for predictions because the random window dataset is not a good.

Stride

Make sure that we use a stride of 1 windows_size when creating the prediction dataset. If there is any overlap, then the results come out really weird. I suspect that it is because I am doing some rescaling without knowing it. Perhaps @AdeelH might have some suggestions for this.

ConcatDataset

We have to do the images 1 at a time. In other words, no ConcatDataset because we need access to the windows attributes. So we can simply loop through the individual datasets.

ds: ConcatDataset = ...

for ids in ds:
   # do apply the same code as above...

Rogue Rescaling

Be careful with the rogue rescaling (as alluded to above). So after going through the above process (using the same dataset for predictions), the final output image was rescaled between [0.0, 1.0]. However, multiplying it by 255/2 solved the problem and we find that it was the same.

# extract image predictions
img_pred: np.ndarray = ds_labels.get_score_arr(ds_i.scene.extent)
# scale it
img_pred *= 255.
# compare
np.testing.assert_array_equal(img_ori, img_pred)

[AdeelH]

Make sure that we use a stride of 1 when creating the prediction dataset. If there is any overlap, then the results come out really weird.

I think a stride of 1 pixel is too small and would result in the maximum overlap. Do you mean stride = 1x chip_sz instead?

If there is any overlap, then the results come out really weird. I suspect that it is because I am doing some rescaling without knowing it. Perhaps @AdeelH might have some suggestions for this.

I think I know what's happening here. Instead of maintaining an average of the predictions for each pixel, SemanticSegmentationSmoothLabels maintains the sum of the values and the number of hits for each pixel serparately, so that it can compute the average when needed. If you access the scores via ds_labels.get_score_arr() instead of ds_labels.pixel_scores directly, it should return the averaged values instead.

[jejjohnson]

I think a stride of 1 pixel is too small and would result in the maximum overlap. Do you mean stride = 1x chip_sz instead?

Yes, sorry, that's exactly what I meant in terms of stride! I've updated the response.

I think I know what's happening here. Instead of maintaining an average of the predictions for each pixel, SemanticSegmentationSmoothLabels maintains the sum of the values and the number of hits for each pixel serparately, so that it can compute the average when needed. If you access the scores via ds_labels.get_score_arr() instead of ds_labels.pixel_scores directly, it should return the averaged values instead.

Yup, that's exactly what I wanted! I'll also add that the ds_labels.get_score_arr() requires a window: Box object. So we just need update this. I'll change this within the text above.

img_pred: np.ndarray = ds_labels.get_score_arr(ds_i.scene.extent)

Thank you for your help!

[AdeelH]

Great!

Just one other note: you should be using ds_i.scene.extent instead of ds_i.scene.bbox in your code above. In your case, both are the same, so it doesn't make a difference, but, in general, they can be different. bbox is used to specify a "crop" of the full image e.g. if bbox=Box(ymin=10, xmin=10, ymax=90, xmax=90) would represent the center 80x80 crop of a 100x100 image. On the other hand, extent in RV parlance is always going to be a Box(0, 0, H, W) box where H and W are the height and width of the bbox.

[jejjohnson]

Ah, great catch! Thank you for the suggestion. I'll update the pseudocode in the previous replies.