Refactor create_polys and to shapefile pipeline.

Author: pounde

handler.py

@@ -1,36 +1,29 @@
 import cv2
 import timeit
-import glob
 import argparse
 import os
 import multiprocessing as mp
 mp.set_start_method('spawn', force=True)
-import sys
-import resource
 import numpy as np
 import raster_processing
-import to_agol
+from utils import to_shapefile
+from utils import to_agol
+from utils import features
 import rasterio.warp
 import torch
 #import ray
 from collections import defaultdict
 from os import makedirs, path
-from functools import partial
 from pathlib import Path
-from shapely.geometry import mapping
 from torch.utils.data import DataLoader
-from yacs.config import CfgNode
 from skimage.morphology import square, dilation
 from tqdm import tqdm
 
 
 from dataset import XViewDataset
 from models import XViewFirstPlaceLocModel, XViewFirstPlaceClsModel
 
-import functools
 import logging
-import struct
-import sys
 
 logger = logging.getLogger()
 
@@ -541,10 +534,11 @@ def main():
     # Get files for creating shapefile and/or pushing to AGOL
     if args.create_shapefile or agol_push.get('dmg'):
         dmg_files = get_files(Path(args.output_directory) / 'dmg')
+        polygons = features.create_polys(dmg_files)
 
     if args.create_shapefile:
         print('Creating shapefile')
-        raster_processing.create_shapefile(dmg_files,
+        to_shapefile.create_shapefile(polygons,
                          Path(args.output_directory).joinpath('shapes') / 'damage.shp',
                          args.destination_crs)
 

raster_processing.py

@@ -218,42 +218,3 @@ def get_tiles(ds, width, height):
             chips.append(outpath.resolve())
 
     return chips
-
-
-def create_shapefile(in_files, out_shapefile, dest_crs):
-
-    """
-    Create shapefiles from input dmg files
-    :param in_files: Files to process for shapefile.
-    :param out_shapefile: File to write shapefile.
-    :param dest_crs: Destination crs.
-    :return: None
-    """
-
-    polygons = []
-    for idx, f in enumerate(in_files):
-        src = rasterio.open(f)
-        crs = src.crs
-        transform = src.transform
-
-        bnd = src.read(1)
-        polys = list(shapes(bnd, transform=transform))
-
-        for geom, val in polys:
-            if val == 0:
-                continue
-            polygons.append((Polygon(shape(geom)), val))
-
-    shp_schema = {
-        'geometry': 'Polygon',
-        'properties': {'dmg': 'int'}
-    }
-
-    # Write out all the multipolygons to the same file
-    with fiona.open(out_shapefile, 'w', 'ESRI Shapefile', shp_schema,
-                    dest_crs) as shp:
-        for polygon, px_val in polygons:
-            shp.write({
-                'geometry': mapping(polygon),
-                'properties': {'dmg': int(px_val)}
-            })

tests/__init__.py

@@ -1,26 +1,26 @@
-import unittest
+from unittest import TestCase
 import handler
 
 
-class TestGetFiles(unittest.TestCase):
+class TestGetFiles(TestCase):
 
     def test_get_files(self):
         self.path = 'data/input/pre'
         self.result = handler.get_files(self.path)
         self.assertEqual(4, len(self.result))
 
 
-class TestReprojectionHelper(unittest.TestCase):
+class TestReprojectionHelper(TestCase):
 
     pass
 
 
-class TestPostprocessAndWrite(unittest.TestCase):
+class TestPostprocessAndWrite(TestCase):
 
     pass
 
 
-class TestFilesClass(unittest.TestCase):
+class TestFilesClass(TestCase):
 
     pass
 

tests/test_handler.py

@@ -1,9 +1,9 @@
-import unittest
+from unittest import TestCase
 
 import inference
 
 
-class TestInference(unittest.TestCase):
+class TestInference(TestCase):
 
     def setUp(self):
         self.opts = inference.Options('sample_data/input/pre')

tests/test_inference.py

@@ -1,4 +1,4 @@
-import unittest
+from unittest import TestCase
 import os
 import rasterio
 from pathlib import Path
@@ -7,7 +7,7 @@
 import handler
 
 
-class TestReproject(unittest.TestCase):
+class TestReproject(TestCase):
 
     def test_reproject(self):
         self.in_file = Path('data/input/pre/tile_337-10160.tif')
@@ -17,7 +17,7 @@ def test_reproject(self):
         self.assertEqual('EPSG:4326', self.test)
 
 
-class TestCreateMosaic(unittest.TestCase):
+class TestCreateMosaic(TestCase):
 
     def setUp(self):
         self.files = handler.get_files(Path('data/input/pre'))
@@ -41,7 +41,7 @@ def test_correct_extent(self):
         self.assertEqual((366642.60000000003, 4104511.1999999997), self.src.transform * (0, 0))
 
 
-class TestGetIntersect(unittest.TestCase):
+class TestGetIntersect(TestCase):
 
     def test_intersect_extent(self):
         self.test = raster_processing.get_intersect(
@@ -69,10 +69,7 @@ def test_intersect_extent(self):
 #         self.result = raster_processing.get_intersect_win(self.mosaic, self.intersect)
 #         self.assertEqual(self.test, self.result)
 
-class TestCreateShapefile(unittest.TestCase):
 
-    pass
-
-class TestCheckDims(unittest.TestCase):
+class TestCheckDims(TestCase):
 
     pass

tests/test_raster_processing.py

@@ -1,2 +1 @@
-import unittest
-import to_agol
+

tests/test_to_agol.py

@@ -0,0 +1,13 @@
+from utils import to_shapefile
+from unittest import TestCase
+from utils.features import create_polys
+from handler import get_files
+from pathlib import Path
+
+
+class Test(TestCase):
+
+    def test_create_shapefile(self):
+        self.files = get_files(Path('data/output/dmg'))
+        self.polys = create_polys(self.files)
+        self.shapefile = to_shapefile(self.polys, Path('~/Downloads'), 'EPSG:4326')

tests/test_to_shapefile.py

@@ -0,0 +1,27 @@
+import rasterio
+from rasterio.features import shapes
+from shapely.geometry import Polygon, shape
+
+def create_polys(in_files):
+
+    """
+    Create palygons to use for feature creation.
+    :param in_files: DMG files to create polygons from.
+    :return: Polygons from dmg files.
+    """
+
+    polygons = []
+    for idx, f in enumerate(in_files):
+        src = rasterio.open(f)
+        crs = src.crs
+        transform = src.transform
+
+        bnd = src.read(1)
+        polys = list(shapes(bnd, transform=transform))
+
+        for geom, val in polys:
+            if val == 0:
+                continue
+            polygons.append((Polygon(shape(geom)), val))
+
+    return polygons

utils/__init__.py

@@ -1,9 +1,6 @@
 import arcgis
-import rasterio
-import shapely.geometry
 from tqdm import tqdm
-from rasterio.features import shapes
-from shapely.geometry import Polygon, shape, MultiPolygon
+from shapely.geometry import MultiPolygon
 
 
 
@@ -45,31 +42,6 @@ def agol_arg_check(args):
     return True
 
 
-def create_polys(in_files):
-    # TODO: Combine this with raster_processing.create_shapefile polygon creation.
-    """
-    Create palygons to use for feature creation.
-    :param in_files: DMG files to create polygons from.
-    :return: Polygons from dmg files.
-    """
-
-    polygons = []
-    for idx, f in enumerate(in_files):
-        src = rasterio.open(f)
-        crs = src.crs
-        transform = src.transform
-
-        bnd = src.read(1)
-        polys = list(shapes(bnd, transform=transform))
-
-        for geom, val in polys:
-            if val == 0:
-                continue
-            polygons.append((Polygon(shape(geom)), val))
-
-    return polygons
-
-
 def create_aoi_poly(features):
 
     """

utils/features.py

@@ -0,0 +1,26 @@
+import fiona
+from shapely.geometry import mapping
+
+def create_shapefile(polygons, out_shapefile, dest_crs):
+
+    """
+    Create shapefile from input polygons
+    :param polygons: Polygons to export as shapefile.
+    :param out_shapefile: Destination shapefile.
+    :param dest_crs: Destination CRS.
+    :return: None
+    """
+
+    shp_schema = {
+            'geometry': 'Polygon',
+            'properties': {'dmg': 'int'}
+        }
+
+    # Write out all the multipolygons to the same file
+    with fiona.open(out_shapefile, 'w', 'ESRI Shapefile', shp_schema,
+                    dest_crs) as shp:
+        for polygon, px_val in polygons:
+            shp.write({
+                'geometry': mapping(polygon),
+                'properties': {'dmg': int(px_val)}
+            })