1# Emulates GDAL's gdal_polygonize.py
2
3import argparse
4import logging
5import subprocess
6import sys
7
8import fiona
9import rasterio
10from rasterio.features import shapes
11
12
13logging.basicConfig(stream=sys.stderr, level=logging.INFO)
14logger = logging.getLogger('rasterio_polygonize')
15
16
17def main(raster_file, vector_file, driver, mask_value):
18
19    with rasterio.Env():
20
21        with rasterio.open(raster_file) as src:
22            image = src.read(1)
23
24        if mask_value is not None:
25            mask = image == mask_value
26        else:
27            mask = None
28
29        results = (
30            {'properties': {'raster_val': v}, 'geometry': s}
31            for i, (s, v)
32            in enumerate(
33                shapes(image, mask=mask, transform=src.transform)))
34
35        with fiona.open(
36                vector_file, 'w',
37                driver=driver,
38                crs=src.crs,
39                schema={'properties': [('raster_val', 'int')],
40                        'geometry': 'Polygon'}) as dst:
41            dst.writerecords(results)
42
43    return dst.name
44
45if __name__ == '__main__':
46
47    parser = argparse.ArgumentParser(
48        description="Writes shapes of raster features to a vector file")
49    parser.add_argument(
50        'input',
51        metavar='INPUT',
52        help="Input file name")
53    parser.add_argument(
54        'output',
55        metavar='OUTPUT',
56        help="Output file name")
57    parser.add_argument(
58        '--output-driver',
59        metavar='OUTPUT DRIVER',
60        help="Output vector driver name")
61    parser.add_argument(
62        '--mask-value',
63        default=None,
64        type=int,
65        metavar='MASK VALUE',
66        help="Value to mask")
67    args = parser.parse_args()
68
69    name = main(args.input, args.output, args.output_driver, args.mask_value)
70
71    print(subprocess.check_output(
72            ['ogrinfo', '-so', args.output, name]))
73