Mask CDL developed pixels that aren't roads or buildings but are clos…

…e to roads

.dockerignore

@@ -6,6 +6,7 @@ county/*
 naip/*
 predictions/*
 road_annotations/*
+road_annotations_for_mask/*
 roads/*
 sample_images/*
 saved_models/*

.gitignore

@@ -8,6 +8,7 @@ naip/*
 predictions/*
 roads/*
 road_annotations/*
+road_annotations_for_mask/*
 sample_images/*
 saved_models/*
 tl_2018_us_state/*

src/annotate_naip_scenes.py

@@ -22,8 +22,10 @@
     COUNTY_FILE,
     NAIP_DIR,
     ROAD_ANNOTATION_DIR,
+    ROAD_ANNOTATION_FOR_MASK_DIR,
     ROAD_BUFFER_METERS,
     ROAD_BUFFER_METERS_DEFAULT,
+    ROAD_BUFFER_METERS_MASK,
     ROAD_DIR,
     ROAD_FORMAT,
 )
@@ -113,17 +115,9 @@ def save_cdl_annotation_for_naip_raster(cdl, naip_file, naip):
     x_cdl, y_cdl = pyproj.transform(proj_naip, proj_cdl, x_naip, y_naip)
 
     cdl_values = get_raster_values(cdl, x_cdl, y_cdl)
-
-    # TODO Mask CDL developed when close to NAIP roads
-
     cdl_values = cdl_values.reshape((naip.meta["height"], naip.meta["width"]))
 
-    profile = naip.profile.copy()
-
-    # Note: the output has the same width, height, and transform as the NAIP raster,
-    # but contains a single band of CDL codes (whereas the NAIP raster contains 4 bands)
-    profile["dtype"] = cdl.profile["dtype"]
-    profile["count"] = 1
+    profile = get_raster_profile(naip, n_bands=1, dtype=cdl.profile["dtype"])
 
     print(f"writing {output_path}")
 
@@ -174,12 +168,7 @@ def save_building_annotation_for_naip_raster(naip_file, naip):
             (naip.meta["height"], naip.meta["width"]), dtype="uint8"
         )
 
-    profile = naip.profile.copy()
-
-    # Note: the output has the same width, height, and transform as the NAIP raster,
-    # but contains a single band of building indicators (whereas the NAIP raster contains 4 bands)
-    profile["dtype"] = "uint8"
-    profile["count"] = 1
+    profile = get_raster_profile(naip, n_bands=1, dtype="uint8")
 
     print(f"writing {output_path}")
     with rasterio.open(output_path, "w", **profile) as output:
@@ -191,12 +180,18 @@ def save_road_annotation_for_naip_raster(counties, naip_file, naip):
     # Note: the road annotation for a given naip_file has the same file name,
     # but is saved to a different directory
     output_path = os.path.join(ROAD_ANNOTATION_DIR, naip_file)
+    output_path_for_mask = os.path.join(ROAD_ANNOTATION_FOR_MASK_DIR, naip_file)
 
-    if os.path.exists(output_path):
-        print(f"{output_path} already exists, skipping")
+    if os.path.exists(output_path) and os.path.exists(output_path_for_mask):
+        print(f"{output_path} and {output_path_for_mask} already exist, skipping")
         return
 
+    # Note: we save two road annotation rasters:
+    # one is used for labeling pixels as roads (for this raster, a small buffer is applied to the road linefiles),
+    # while the second annotation raster is used to mask CDL developed pixels that are within a large buffer
+    # (but not a small buffer) of the road linefiles
     road_geometries = []
+    road_geometries_for_mask = []
 
     for county in counties:
 
@@ -228,6 +223,8 @@ def save_road_annotation_for_naip_raster(counties, naip_file, naip):
             )
             road_geometries.append(road_geometry_transformed.buffer(road_buffer_meters))
 
+            road_geometries_for_mask.append(road_geometry_transformed.buffer(ROAD_BUFFER_METERS_MASK))
+
     road_values = rasterize(
         road_geometries,
         out_shape=(naip.meta["height"], naip.meta["width"]),
@@ -236,22 +233,40 @@ def save_road_annotation_for_naip_raster(counties, naip_file, naip):
         dtype="uint8",
     )
 
-    # TODO Copied from CDL code, put in a function
-    profile = naip.profile.copy()
-
-    # Note: the output has the same width, height, and transform as the NAIP raster,
-    # but contains a single band of road indicators (whereas the NAIP raster contains 4 bands)
-    profile["dtype"] = "uint8"
-    profile["count"] = 1
+    road_values_for_mask = rasterize(
+        road_geometries_for_mask,
+        out_shape=(naip.meta["height"], naip.meta["width"]),
+        transform=naip.transform,
+        all_touched=True,
+        dtype="uint8",
+    )
 
-    print(f"writing {output_path}")
+    profile = get_raster_profile(naip, n_bands=1, dtype="uint8")
 
-    if not os.path.exists(ROAD_ANNOTATION_DIR):
-        os.makedirs(ROAD_ANNOTATION_DIR)
+    for road_annotation_dir in [ROAD_ANNOTATION_FOR_MASK_DIR, ROAD_ANNOTATION_DIR]:
+        if not os.path.exists(road_annotation_dir):
+            os.makedirs(road_annotation_dir)
 
+    print(f"writing {output_path}")
     with rasterio.open(output_path, "w", **profile) as output:
         output.write(road_values.astype(profile["dtype"]), 1)
 
+    print(f"writing {output_path_for_mask}")
+    with rasterio.open(output_path_for_mask, "w", **profile) as output:
+        output.write(road_values_for_mask.astype(profile["dtype"]), 1)
+
+
+def get_raster_profile(naip, n_bands, dtype):
+
+    profile = naip.profile.copy()
+
+    # Note: the output has the same width, height, and transform as the NAIP raster,
+    # but contains n_bands bands (whereas the NAIP raster contains 4 bands)
+    profile["dtype"] = dtype
+    profile["count"] = n_bands
+
+    return profile
+
 
 def get_counties(raster):
 

src/constants.py

@@ -25,6 +25,7 @@
 COUNTY_FILE = "tl_2017_us_county.shp"
 
 ROAD_ANNOTATION_DIR = "./road_annotations"
+ROAD_ANNOTATION_FOR_MASK_DIR = "./road_annotations_for_mask"
 ROAD_DIR = "./roads"
 ROAD_FORMAT = "tl_2017_{county}_roads.shp"
 
@@ -33,6 +34,9 @@
 ROAD_BUFFER_METERS_DEFAULT = 3.0
 ROAD_BUFFER_METERS = {"S1100": 5.0, "S1200": 5.0}
 
+# Note: roads are buffered by this amount in order to implement CDL developed masking logic
+ROAD_BUFFER_METERS_MASK = 25.0
+
 HAS_BUILDINGS = "has_buildings"
 HAS_ROADS = "has_roads"
 IS_MAJORITY_FOREST = "is_majority_forest"

src/fit_model.py

@@ -29,6 +29,7 @@
     MODEL_CONFIG,
     NAIP_DIR,
     ROAD_ANNOTATION_DIR,
+    ROAD_ANNOTATION_FOR_MASK_DIR,
     SAVED_MODELS_DIR,
 )
 from generator import get_generator
@@ -55,22 +56,27 @@ def recode_cdl_values(cdl_values, cdl_mapping, label_encoder):
     return cdl_recoded
 
 
-def get_y_combined(y_cdl_recoded, y_road, y_building, label_encoder):
+def get_y_combined(y_cdl_recoded, y_road, y_road_for_mask, y_building, label_encoder):
 
     ## This function "burns" roads and buildings into the recoded CDL raster
     ## The road and building datasets appear to be more reliable than the CDL,
     ## and they are higher resolution, so roads and buildings take precedence over CDL codes
 
-    # TODO Make sure y_cdl_recoded, y_road and y_building get garbage collected!
-    # TODO Make sure y_combined is uint8
     y_combined = y_cdl_recoded.copy()
 
     road = label_encoder.transform(["road"])[0]
     y_combined[np.where(y_road)] = road
 
+    # Note: this could in theory overwrite some roads
     building = label_encoder.transform(["building"])[0]
     y_combined[np.where(y_building)] = building
 
+    # Note: this masks out pixels that are (1) CDL developed,
+    # (2) not roads or buildings, and (3) close to a road (without being covered by a road)
+    mask = label_encoder.transform(["mask"])[0]
+    developed = label_encoder.transform(["developed"])[0]
+    y_combined[np.where(np.logical_and(y_road_for_mask, y_combined == developed))] = mask
+
     return y_combined
 
 
@@ -102,11 +108,16 @@ def get_annotated_scenes(naip_paths, label_encoder, cdl_mapping):
 
             y_road = road_annotation.read()
 
+        road_annotation_for_mask_path = os.path.join(ROAD_ANNOTATION_FOR_MASK_DIR, naip_file)
+        with rasterio.open(road_annotation_for_mask_path) as road_annotation_for_mask:
+
+            y_road_for_mask = road_annotation_for_mask.read()
+
         building_annotation_path = os.path.join(BUILDING_ANNOTATION_DIR, naip_file)
         with rasterio.open(building_annotation_path) as building_annotation:
             y_building = building_annotation.read()
 
-        y_combined = get_y_combined(y_cdl_recoded, y_road, y_building, label_encoder)
+        y_combined = get_y_combined(y_cdl_recoded, y_road, y_road_for_mask, y_building, label_encoder)
 
         # Note: swap NAIP and CDL shape from (band, height, width) to (width, height, band)
         annotated_scenes.append([np.swapaxes(X, 0, 2), np.swapaxes(y_combined, 0, 2)])