diff --git a/geofastai/raster.py b/geofastai/raster.py index 6d46304..02a2dfd 100644 --- a/geofastai/raster.py +++ b/geofastai/raster.py @@ -10,13 +10,18 @@ class RasterBase(fai.TensorImage): def create(cls, fname, **kwargs): data = rio.open(fname).read() return cls(data) - + + def band_quantile(self, q): + return self.reshape(self.shape[0], -1, 1).quantile(q, axis=1, keepdims=True) + class RasterVNIR(RasterBase): - def show(self, ctx=None, bands=[3,2,1], **kwargs): + def show(self, ctx=None, bands=[3,2,1], q=0.001, **kwargs): im = self[bands] im = im.to(fai.torch.float) - im = im / im.quantile(.9) + vmin = im.band_quantile(q) + vmax = im.band_quantile(1 - q) + im = (im - vmin) / (vmax - vmin) im = im.clip(0, 1) return fai.show_image(im, ctx=ctx, vmax=300, **kwargs)