# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os.path as osp
import numpy as np
from typing import List, Tuple, Union
from paddlers.transforms.functions import to_uint8 as raster2uint8
try:
from osgeo import gdal
except:
import gdal
class Raster:
def __init__(self,
path: str,
band_list: Union[List[int], Tuple[int], None]=None,
to_uint8: bool=False) -> None:
""" Class of read raster.
Args:
path (str): The path of raster.
band_list (Union[List[int], Tuple[int], None], optional):
band list (start with 1) or None (all of bands). Defaults to None.
to_uint8 (bool, optional):
Convert uint8 or return raw data. Defaults to False.
"""
super(Raster, self).__init__()
if osp.exists(path):
self.path = path
self.ext_type = path.split(".")[-1]
if self.ext_type.lower() in ["npy", "npz"]:
self._src_data = None
else:
try:
# raster format support in GDAL:
# https://www.osgeo.cn/gdal/drivers/raster/index.html
self._src_data = gdal.Open(path)
except:
raise TypeError("Unsupported data format: `{}`".format(self.ext_type))
self._getInfo()
self.to_uint8 = to_uint8
self.setBands(band_list)
else:
raise ValueError("The path {0} not exists.".format(path))
def setBands(self,
band_list: Union[List[int], Tuple[int], None]) -> None:
""" Set band of data.
Args:
band_list (Union[List[int], Tuple[int], None]):
band list (start with 1) or None (all of bands).
"""
if band_list is not None:
if len(band_list) > self.bands:
raise ValueError("The lenght of band_list must be less than {0}.".format(str(self.bands)))
if max(band_list) > self.bands or min(band_list) < 1:
raise ValueError("The range of band_list must within [1, {0}].".format(str(self.bands)))
self.band_list = band_list
def getArray(self,
start_loc: Union[List[int], Tuple[int], None]=None,
block_size: Union[List[int], Tuple[int]]=[512, 512]) -> np.ndarray:
""" Get ndarray data
Args:
start_loc (Union[List[int], Tuple[int], None], optional):
Coordinates of the upper left corner of the block, if None means return full image.
block_size (Union[List[int], Tuple[int]], optional):
Block size. Defaults to [512, 512].
Returns:
np.ndarray: data's ndarray.
"""
if self._src_data is not None:
if start_loc is None:
return self._getAarray()
else:
return self._getBlock(start_loc, block_size)
else:
print("Numpy doesn't support blocking temporarily.")
return self._getNumpy()
def _getInfo(self) -> None:
if self._src_data is not None:
self.bands = self._src_data.RasterCount
self.width = self._src_data.RasterXSize
self.height = self._src_data.RasterYSize
self.geot = self._src_data.GetGeoTransform()
self.proj = self._src_data.GetProjection()
else:
d_shape = self._getNumpy().shape
if len(d_shape) == 3:
self.height, self.width, self.bands = d_shape
else:
self.height, self.width = d_shape
self.bands = 1
self.geot = None
self.proj = None
def _getNumpy(self):
ima = np.load(self.path)
if self.band_list is not None:
band_array = []
for b in self.band_list:
band_i = ima[:, :, b - 1]
band_array.append(band_i)
ima = np.stack(band_array, axis=0)
return ima
def _getAarray(self, window: Union[None, List[int], Tuple[int]]=None) -> np.ndarray:
if window is not None:
xoff, yoff, xsize, ysize = window
if self.band_list is None:
if window is None:
ima = self._src_data.ReadAsArray()
else:
ima = self._src_data.ReadAsArray(xoff, yoff, xsize, ysize)
else:
band_array = []
for b in self.band_list:
if window is None:
band_i = self._src_data.GetRasterBand(b).ReadAsArray()
else:
band_i = self._src_data.GetRasterBand(b).ReadAsArray(xoff, yoff, xsize, ysize)
band_array.append(band_i)
ima = np.stack(band_array, axis=0)
if self.bands == 1:
# the type is complex means this is a SAR data
if isinstance(type(ima[0, 0]), complex):
ima = abs(ima)
else:
ima = ima.transpose((1, 2, 0))
if self.to_uint8 is True:
ima = raster2uint8(ima)
return ima
def _getBlock(self,
start_loc: Union[List[int], Tuple[int]],
block_size: Union[List[int], Tuple[int]]=[512, 512]) -> np.ndarray:
if len(start_loc) != 2 or len(block_size) != 2:
raise ValueError("The length start_loc/block_size must be 2.")
xoff, yoff = start_loc
xsize, ysize = block_size
if (xoff < 0 or xoff > self.width) or (yoff < 0 or yoff > self.height):
raise ValueError(
"start_loc must be within [0-{0}, 0-{1}].".format(str(self.width), str(self.height)))
if xoff + xsize > self.width:
xsize = self.width - xoff
if yoff + ysize > self.height:
ysize = self.height - yoff
ima = self._getAarray([int(xoff), int(yoff), int(xsize), int(ysize)])
h, w = ima.shape[:2] if len(ima.shape) == 3 else ima.shape
if self.bands != 1:
tmp = np.zeros((block_size[0], block_size[1], self.bands), dtype=ima.dtype)
tmp[:h, :w, :] = ima
else:
tmp = np.zeros((block_size[0], block_size[1]), dtype=ima.dtype)
tmp[:h, :w] = ima
return tmp