Repository: rasterio/rasterio

Issue:
rasterize to allow data types other than uint8
Currently the _rasterize function has the data type hardcoded to bytes. This won't work for many cases of continuous value rasters (most of data I work with is float32 or uint16). Is there a way to pass the datatype in as an argument while keeping the byte default?



Discussion:
For the benefit of future google searches, this is the error:

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-13-a716cdec5847> in <module>()
      3 image = features.rasterize(
      4             ((feat.geom, feat.cropval) for i, feat in df.iterrows()),
----> 5             out_shape=(600,600))

/home/mperry/env/geopandas/lib/python2.7/site-packages/rasterio/features.py in rasterize(shapes, out_shape, fill, output, transform, all_touched)
     87                 raise ValueError(
     88                     "Shape number %i, value '%s' is not uint8/ubyte" % (
---> 89                         index, value))
     90             geoms.append((geometry, value))
     91         else:

ValueError: Shape number 0, value '0.0' is not uint8/ubyte
```

I don't see any reason why not. The inverse operation (`shapes()`) on the other hand, is probably only feasible for uint8.

Actually, `shapes()` can also be adapted to use use a 32bit float buffer, via [GDALFPolygonize](http://www.gdal.org/gdal__alg_8h.html#a7a789015334d677afcbef67e5a6b4a7c)

float32 should be straightforward (see [GDALRasterizeGeometries](http://www.gdal.org/gdal__alg_8h.html#a50caf4bc34703f0bcf515ecbe5061a0a): "output raster may be of any GDAL supported datatype, though currently internally the burning is done either as GDT_Byte or GDT_Float32."  Presumably uint16 would be handled OK internally as a GDT_Float32, would need testing.

within _features.pyx, it looks like we are upcasting the feature values to doubles anyway, so no issue there other than removing some of the validation upstream of there.  Handling other data types at that level should be as easy as substituting in different IO functions from io.pyx: e.g., io_uint16, io_float32

The bigger issue as I see it is the function signature for features.py::rasterize.  Namely, if output is specified, we could use that to determine data type for validation and downstream use.  Otherwise, we'd need to add an additional parameter for data type, and then extra logic to evaluate feature values and default_value against that.  We may also want to add a dictionary of default values by data type to dtypes.py, and use those based on target data type.

@sgillies is there another precedent for passing in data type as a parameter?

@mwtoews true... i was just thinking of what the shapes of non-thematic rasters would look like, every pixel its own polygon seems like a common outcome.

@brendan-ward precedent in Rasterio? No, but I think it would make sense to follow the params of a Numpy function like numpy.zeros or [numpy.trace](http://docs.scipy.org/doc/numpy/reference/generated/numpy.trace.html#numpy.trace) and add a dtype=None argument.

@sgillies - re: the dtype in `shapes()`, the case of "each pixel is a polygon" could certainly happen if you supplied a continuous value float. But that's up to the user IMO. A more common workflow would be to classify the continuous values into discrete categories resulting in a UInt8/UInt16 array then run `shapes()` on that. But we'd still want support for other data types, at least UInt16.

I'm :thumbsup: :thumbsup: for generalizing this. Will be an excellent enhancement. Who's going to pick it up? I'm completely slammed through 9/12.

I'm completely slammed too, but I think I've got a path forward on this.  The code changes involved are not hard.  Can you give me a week to try and either get this done or bow out?

And to be clear - I was only volunteering for the `rasterize` data types, not `shapes`

@brendan-ward go for it. I made a separate issues for `shapes()`.

@perrygeo Are you OK with data type being auto-detected based on feature values, unless you specifically pass in a value for `dtype`?

I'm thinking the function signature could be

```
rasterize(
        shapes, 
        out_shape=None, 
        fill=0, 
        output=None,
        transform=IDENTITY,
        all_touched=False,
        default_value=None,
        dtype=None)
```

and we handle `dtype` like numpy, and autodetect the type if `dtype is None` and use that for the resultant image array.  We'll have to do a little validation there; not all numpy dtypes are valid for this function.

That sounds great, very flexible but with intelligent defaults. 

Heads up: some data type changes landed in #141. Good news is that the changes make things a little easier for users.

I've been making good progress on this, but have a couple topics for discussion:
- use 1 as default value for all data types (right now it is 255): after digging in more, it felt less appropriate to use the upper value of a given data type as it's default, esp. when precision may vary a little bit across architectures for things like float64.  1 seems like a good value to indicate the location of features, when 0 is the default background.  And really, the case where data type matters is when user specified default or feature value range needs a particular data type, so a default of 1 doesn't interfere.
- modify RasterUpdater to work w/ memory driver: right now, several places in rasterio use boilerplate to create a memory dataset (of a given type), get the first band, and use that for IO with numpy arrays.    After a brief foray into a dedicated memory based class, it looks like simply updating RasterUpdater with the appropriate logic would greatly consolidate that boilerplate down.  However, this would change some of the patterns when using a memory dataset; it couldn't auto-determine things by reading from the file (e.g., CRS), so we'd have to insist that those be passed in for memory drivers (if required).

@brendan-ward :+1: on changing the default value. I'm also fine with refactoring and using RasterUpdater internally so long as the memory driver doesn't leak into the user interface. I think there ought to be one obvious way to have rasters in memory (numpy arrays) and not two different ways (numpy arrays and GDAL memory datasets).

@sgillies I was thinking that the user would pass the "MEM" driver to the call to rasterio.open - this sounds like it conflicts with your desire to not have it in the user interface.  What level are you considering user interface?

I agree about primary way for working w/ memory arrays (numpy).  This new functionality is more useful to help w/ IO w/ GDAL functions (meaning we really only need this accessible in the cython stuff).  I was hoping to leverage in the benefits of the "with" usage pattern for rasterio.open, but maybe that doesn't make as much sense if we are only using from cython.  

Are you open to a lighter weight, dedicated class we use just from cython to handle memory datasets for IO w/ GDAL?

A dedicated class seems exactly right. Let's do that.

Also, heads up: since https://github.com/mapbox/rasterio/commit/4e53e3ca2d448e1abc3ad7522b3ac2292eacfb95 we have a new DatasetReader base class that RasterReader/Updater derive from (in case you're looking for where all the methods went).

@sgillies the generator you added to features.py in b60e4635 makes it harder than I'd like to iterate over all values in the input shapes, put them in an array, and determine the type from the aggregate.  Are you tied to this being a generator here?

I was thinking I could rewrite that operation as one that does more of a filter operation (discarding bad shapes & logging exceptions) while adding their values to an array and using some bounds checking logic on the min and max of that array to determine the dtype to use.

I've got a first pass at the dedicated class, but I'm not sure that I totally like the API.  

Given a numpy 2D array, it constructs an in-memory dataset with 1 band for IO with GDAL functions.  It keeps a couple attributes for passing into those GDAL functions (e.g., self.dataset, self.band_ids, etc).  It writes the initial values of that array to the in-memory band, and uses the `with` pattern to close the dataset later.

Thoughts?

(in _io.pyx)

```
cdef class InMemoryRaster:
    def __init__(self, image, transform):
        self.dataset = NULL

        cdef void *memdriver = _gdal.GDALGetDriverByName("MEM")

        self.band_ids[0] = 1

        self.dataset = _gdal.GDALCreate(
            memdriver,
            "output",
            image.shape[1],
            image.shape[0],
            1,
            <_gdal.GDALDataType>dtypes.dtype_rev[image.dtype.name],
            NULL
        )

        if self.dataset == NULL:
            raise ValueError("NULL output datasource")

        for i in range(6):
            self.transform[i] = transform[i]
        err = _gdal.GDALSetGeoTransform(self.dataset, self.transform)
        if err:
            raise ValueError("transform not set: %s" % transform)

        self.band = _gdal.GDALGetRasterBand(self.dataset, 1)
        if self.band == NULL:
            raise ValueError("NULL output band: {0}".format(i))

        self.write(image)

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        if self.dataset != NULL:
            _gdal.GDALClose(self.dataset)

    def read(self, image):
        return io_auto(image, self.band, False)

    def write(self, image):
        return io_auto(image, self.band, True)
```

And usage (e.g., in _features.pyx):

```
        with InMemoryRaster(image, transform) as mem:
            retval = _gdal.GDALRasterizeGeometries(
                        mem.dataset, 1, mem.band_ids,
                        num_geometries, ogr_geoms,
                        NULL, mem.transform, pixel_values,
                        options, NULL, NULL)

            # Write the in-memory band back to the image.
            mem.read(image)
```

Note: `io_auto` above routes to the appropriate io function for the dtype of the image.  I could encapsulate that here only, but it seemed like a nice general utility function.

@brendan-ward seems good to me. As its an implementation detail and not part of the API, we can tweak it as we go. It'll eliminate a lot of redundant code in Rasterio, for sure.

No, I'm not wedded to a generator in features.py. So long as the functions take and return iterators, I'm okay with internal changes.


Resolving patch:
diff --git a/rasterio/dtypes.py b/rasterio/dtypes.py
--- a/rasterio/dtypes.py
+++ b/rasterio/dtypes.py
@@ -71,3 +71,28 @@ def check_dtype(dt):
         except:
             return False
     return True
+
+
+def get_minimum_int_dtype(values):
+    """
+    Uses range checking to determine the minimum integer data type required
+    to represent values.
+
+    :param values: numpy array
+    :return: named data type that can be later used to create a numpy dtype
+    """
+
+    min_value = values.min()
+    max_value = values.max()
+    
+    if min_value >= 0:
+        if max_value <= 255:
+            return uint8
+        elif max_value <= 65535:
+            return uint16
+        elif max_value <= 4294967295:
+            return uint32
+    elif min_value >= -32768 and max_value <= 32767:
+        return int16
+    elif min_value >= -2147483648 and max_value <= 2147483647:
+        return int32
diff --git a/rasterio/features.py b/rasterio/features.py
--- a/rasterio/features.py
+++ b/rasterio/features.py
@@ -10,6 +10,7 @@
 import rasterio
 from rasterio._features import _shapes, _sieve, _rasterize
 from rasterio.transform import IDENTITY, guard_transform
+from rasterio.dtypes import get_minimum_int_dtype
 
 
 log = logging.getLogger('rasterio')
@@ -68,10 +69,13 @@ def sieve(image, size, connectivity=4, output=None):
 
 def rasterize(
         shapes, 
-        out_shape=None, fill=0, output=None,
+        out_shape=None,
+        fill=0,
+        output=None,
         transform=IDENTITY,
         all_touched=False,
-        default_value=255):
+        default_value=1,
+        dtype=None):
     """Returns an image array with points, lines, or polygons burned in.
 
     A different value may be specified for each shape.  The shapes may
@@ -80,9 +84,10 @@ def rasterize(
     of image elements determines whether they are updated, but all
     touched elements may be optionally updated.
 
-    :param shapes: an iterator over Fiona style geometry objects (with
-    a default value of 255) or an iterator over (geometry, value) pairs.
-    Values must be unsigned integer type (uint8).
+    Valid data types are: int16, int32, uint8, uint16, uint32, float32, float64
+
+    :param shapes: an iterator over Fiona style geometry objects (with a default
+    value of default_value) or an iterator over (geometry, value) pairs.
 
     :param transform: GDAL style geotransform to be applied to the
     image.
@@ -93,54 +98,102 @@ def rasterize(
 
     :param all_touched: if True, will rasterize all pixels touched, 
     otherwise will use GDAL default method.
-    :param default_value: value burned in for shapes if not provided as part of shapes.  Must be unsigned integer type (uint8).
+    :param default_value: value burned in for shapes if not provided as part
+    of shapes.
     """
 
-    if not isinstance(default_value, int) or (
-            default_value > 255 or default_value < 0):
-        raise ValueError("default_value %s is not uint8/ubyte" % default_value)
-
-    def shape_source():
-        """A generator that screens out non-geometric objects and does
-        its best to make sure that no NULLs get through to 
-        GDALRasterizeGeometries."""
-        for index, item in enumerate(shapes):
-            try:
-                
-                if isinstance(item, (tuple, list)):
-                    geom, value = item
-                    # TODO: relax this for other data types.
-                    if not isinstance(value, int) or value > 255 or value < 0:
-                        raise ValueError(
-                            "Shape number %i, value '%s' is not uint8/ubyte" % (
-                                index, value))
-                else:
-                    geom = item
-                    value = default_value
-                geom = getattr(geom, '__geo_interface__', None) or geom
-                if (not isinstance(geom, dict) or 
-                    'type' not in geom or 'coordinates' not in geom):
-                    raise ValueError(
-                        "Object %r at index %d is not a geometric object" %
-                        (geom, index))
-                yield geom, value
-            except Exception:
-                log.exception("Exception caught, skipping shape %d", index)
-
-    if out_shape is not None:
-        out = np.empty(out_shape, dtype=rasterio.ubyte)
-        out.fill(fill)
-    elif output is not None:
-        if np.dtype(output.dtype) != np.dtype(rasterio.ubyte):
-            raise ValueError("Output image must be dtype uint8/ubyte")
-        out = output
+    valid_dtypes = ('int16', 'int32', 'uint8', 'uint16', 'uint32', 'float32',
+                    'float64')
+
+    def get_valid_dtype(values):
+        values_dtype = values.dtype
+        if values_dtype.kind == 'i':
+            values_dtype = np.dtype(get_minimum_int_dtype(values))
+        if values_dtype.name in valid_dtypes:
+            return values_dtype
+        return None
+
+    def can_cast_dtype(values, dtype):
+        if values.dtype.name == np.dtype(dtype).name:
+            return True
+        elif values.dtype.kind == 'f':
+            return np.allclose(values, values.astype(dtype))
+        else:
+            return np.array_equal(values, values.astype(dtype))
+
+    if fill != 0:
+        fill_array = np.array([fill])
+        if get_valid_dtype(fill_array) is None:
+            raise ValueError('fill must be one of these types: %s'
+                             % (', '.join(valid_dtypes)))
+        elif dtype is not None and not can_cast_dtype(fill_array, dtype):
+            raise ValueError('fill value cannot be cast to specified dtype')
+
+
+    if default_value != 1:
+        default_value_array = np.array([default_value])
+        if get_valid_dtype(default_value_array) is None:
+            raise ValueError('default_value must be one of these types: %s'
+                             % (', '.join(valid_dtypes)))
+        elif dtype is not None and not can_cast_dtype(default_value_array,
+                                                      dtype):
+            raise ValueError('default_value cannot be cast to specified dtype')
+
+    valid_shapes = []
+    shape_values = []
+    for index, item in enumerate(shapes):
+        try:
+            if isinstance(item, (tuple, list)):
+                geom, value = item
+            else:
+                geom = item
+                value = default_value
+            geom = getattr(geom, '__geo_interface__', None) or geom
+            if (not isinstance(geom, dict) or
+                'type' not in geom or 'coordinates' not in geom):
+                raise ValueError(
+                    'Object %r at index %d is not a geometry object' %
+                    (geom, index))
+            valid_shapes.append((geom, value))
+            shape_values.append(value)
+        except Exception:
+            log.exception('Exception caught, skipping shape %d', index)
+
+    if not valid_shapes:
+        raise ValueError('No valid shapes found for rasterize.  Shapes must be '
+                         'valid geometry objects')
+
+    shape_values = np.array(shape_values)
+    values_dtype = get_valid_dtype(shape_values)
+    if values_dtype is None:
+        raise ValueError('shape values must be one of these dtypes: %s' %
+                         (', '.join(valid_dtypes)))
+
+    if dtype is None:
+        dtype = values_dtype
+    elif np.dtype(dtype).name not in valid_dtypes:
+        raise ValueError('dtype must be one of: %s' % (', '.join(valid_dtypes)))
+    elif not can_cast_dtype(shape_values, dtype):
+        raise ValueError('shape values could not be cast to specified dtype')
+
+    if output is not None:
+        if np.dtype(output.dtype).name not in valid_dtypes:
+            raise ValueError('Output image dtype must be one of: %s' 
+                             % (', '.join(valid_dtypes)))
+        if not can_cast_dtype(shape_values, output.dtype):
+            raise ValueError('shape values cannot be cast to dtype of output '
+                             'image')
+
+    elif out_shape is not None:
+        output = np.empty(out_shape, dtype=dtype)
+        output.fill(fill)
     else:
-        raise ValueError("An output image must be provided or specified")
+        raise ValueError('Either an output shape or image must be provided')
     
     transform = guard_transform(transform)
 
     with rasterio.drivers():
-        _rasterize(shape_source(), out, transform.to_gdal(), all_touched)
+        _rasterize(valid_shapes, output, transform.to_gdal(), all_touched)
     
-    return out
+    return output
 


Test patch:
diff --git a/tests/test_features_rasterize.py b/tests/test_features_rasterize.py
--- a/tests/test_features_rasterize.py
+++ b/tests/test_features_rasterize.py
@@ -11,28 +11,31 @@
 def test_rasterize_geometries():
     rows = cols = 10
     transform = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0)
-    geometry = {'type':'Polygon','coordinates':[[(2,2),(2,4.25),(4.25,4.25),(4.25,2),(2,2)]]}
+    geometry = {
+        'type':'Polygon',
+        'coordinates':[[(2, 2), (2, 4.25), (4.25, 4.25), (4.25, 2), (2, 2)]]
+    }
 
     with rasterio.drivers():
         # we expect a subset of the pixels using default mode
         result = rasterize([geometry], out_shape=(rows, cols))
         truth = numpy.zeros((rows, cols))
-        truth[2:4, 2:4] = 255
-        assert (result == truth).min() == True
+        truth[2:4, 2:4] = 1
+        assert numpy.array_equal(result, truth)
 
         # we expect all touched pixels
         result = rasterize(
                     [geometry], out_shape=(rows, cols), all_touched=True)
         truth = numpy.zeros((rows, cols))
-        truth[2:5, 2:5] = 255
-        assert (result == truth).min() == True
+        truth[2:5, 2:5] = 1
+        assert numpy.array_equal(result, truth)
 
         # we expect the pixel value to match the one we pass in
         value = 5
         result = rasterize([(geometry, value)], out_shape=(rows, cols))
         truth = numpy.zeros((rows, cols))
         truth[2:4, 2:4] = value
-        assert (result == truth).min() == True
+        assert numpy.array_equal(result, truth)
         
         # Check the fill and default transform.
         # we expect the pixel value to match the one we pass in
@@ -43,7 +46,84 @@ def test_rasterize_geometries():
             fill=1 )
         truth = numpy.ones((rows, cols))
         truth[2:4, 2:4] = value
-        assert (result == truth).min() == True
+        assert numpy.array_equal(result, truth)
+
+
+def test_rasterize_dtype():
+    rows = cols = 10
+    transform = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0)
+    geometry = {
+        'type':'Polygon',
+        'coordinates':[[(2, 2), (2, 4.25), (4.25, 4.25), (4.25, 2), (2, 2)]]
+    }
+
+    with rasterio.drivers():
+        # Supported types should all work properly
+        supported_types = (('int16', -32768), ('int32', -2147483648),
+                     ('uint8', 255), ('uint16', 65535), ('uint32', 4294967295),
+                     ('float32', 1.434532), ('float64', -98332.133422114))
+
+        for dtype, default_value in supported_types:
+            truth = numpy.zeros((rows, cols), dtype=dtype)
+            truth[2:4, 2:4] = default_value
+
+            result = rasterize(
+                [geometry],
+                out_shape=(rows, cols),
+                default_value=default_value,
+                dtype=dtype
+            )
+            assert numpy.array_equal(result, truth)
+            assert numpy.dtype(result.dtype) == numpy.dtype(truth.dtype)
+
+            result = rasterize(
+                [(geometry, default_value)],
+                out_shape=(rows, cols)
+            )
+            if numpy.dtype(dtype).kind == 'f':
+                assert numpy.allclose(result, truth)
+            else:
+                assert numpy.array_equal(result, truth)
+            # Since dtype is auto-detected, it may not match due to upcasting
+
+        # Unsupported types should all raise exceptions
+        unsupported_types = (('int8', -127), ('int64', 20439845334323),
+                             ('float16', -9343.232))
+
+        for dtype, default_value in unsupported_types:
+            with pytest.raises(ValueError):
+                rasterize(
+                    [geometry],
+                    out_shape=(rows, cols),
+                    default_value=default_value,
+                    dtype=dtype
+                )
+
+            with pytest.raises(ValueError):
+                rasterize(
+                    [(geometry, default_value)],
+                    out_shape=(rows, cols),
+                    dtype=dtype
+                )
+
+        # Mismatched values and dtypes should raise exceptions
+        mismatched_types = (('uint8', 3.2423), ('uint8', -2147483648))
+        for dtype, default_value in mismatched_types:
+            print("Testing mismatched dtype {0} with value {1}".format(dtype, default_value))
+            with pytest.raises(ValueError):
+                rasterize(
+                    [geometry],
+                    out_shape=(rows, cols),
+                    default_value=default_value,
+                    dtype=dtype
+                )
+
+            with pytest.raises(ValueError):
+                rasterize(
+                    [(geometry, default_value)],
+                    out_shape=(rows, cols),
+                    dtype=dtype
+                )
 
 
 def test_rasterize_geometries_symmetric():
@@ -55,4 +135,5 @@ def test_rasterize_geometries_symmetric():
     with rasterio.drivers():
         s = shapes(truth, transform=transform)
         result = rasterize(s, out_shape=(rows, cols), transform=transform)
-        assert (result == truth).min() == True
+        assert numpy.array_equal(result, truth)
+