utils overhaul

Author: chrieke

preprocessing.ipynb

@@ -4,5 +4,4 @@
 import utils.geo
 import utils.img
 import utils.coco
-import utils.viz
 

utils/__init__.py

@@ -6,30 +6,34 @@
 import itertools
 
 from shapely.geometry import Polygon, MultiPolygon
-import utils.other
 import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.collections import PatchCollection
+from descartes import PolygonPatch
+from PIL import Image as pilimage
+
+import utils.other
 
 
-def train_test_split_coco(chips_stats: Dict) -> Tuple[List, List]:
-    chips_list = list(chips_stats.keys())
+def train_test_split(chip_dfs: Dict, test_size=0.2) -> Tuple[Dict, Dict]:
+    """Split chips into training and test set"""
+    chips_list = list(chip_dfs.keys())
     random.seed(1)
     random.shuffle(chips_list)
-    split_idx = round(len(chips_list) * 0.2)  # 80% train, 20% test.
+    split_idx = round(len(chips_list) * test_size)
     train_split = chips_list[split_idx:]
     val_split = chips_list[:split_idx]
 
-    # Apply split to geometries/stats.
-    train_chip_dfs = {k: chips_stats[k] for k in sorted(train_split)}
-    val_chip_dfs = {k.replace('train', 'val'): chips_stats[k] for k in sorted(val_split)}
+    train_chip_dfs = {k: chip_dfs[k] for k in sorted(train_split)}
+    val_chip_dfs = {k.replace('train', 'val'): chip_dfs[k] for k in sorted(val_split)}
 
     return train_chip_dfs, val_chip_dfs
 
 
-def format_coco(set_: Dict, chip_width: int, chip_height: int):
-    """
-    Format extracted chip geometries to COCO json format.
+def format_coco(chip_dfs: Dict, chip_width: int, chip_height: int):
+    """Format train and test chip geometries to COCO json format.
 
-    Coco train/val have specific ids, formatting requires the split data..
+    COCO train and val set have specific ids.
     """
     cocojson = {
         "info": {},
@@ -38,7 +42,7 @@ def format_coco(set_: Dict, chip_width: int, chip_height: int):
                         'id': 1,   # id needs to match category_id.
                         'name': 'agfields_singleclass'}]}
 
-    for key_idx, key in enumerate(set_.keys()):
+    for key_idx, key in enumerate(chip_dfs.keys()):
         if 'train' in key:
             chip_id = int(key[21:])
         elif 'val' in key:
@@ -50,16 +54,16 @@ def format_coco(set_: Dict, chip_width: int, chip_height: int):
                       "width": chip_height})
         cocojson.setdefault('images', []).append(key_image)
 
-        for row_idx, row in set_[key]['chip_df'].iterrows():
+        for row_idx, row in chip_dfs[key]['chip_df'].iterrows():
             # Convert geometry to COCO segmentation format:
             # From shapely POLYGON ((x y, x1 y2, ..)) to COCO [[x, y, x1, y1, ..]].
             # The annotations were encoded by RLE, except for crowd region (iscrowd=1)
             coco_xy = list(itertools.chain.from_iterable((x, y) for x, y in zip(*row.geometry.exterior.coords.xy)))
-            coco_xy = [round(xy, 2) for xy in coco_xy]
+            coco_xy = [round(coords, 2) for coords in coco_xy]
             # Add COCO bbox in format [minx, miny, width, height]
             bounds = row.geometry.bounds  # COCO bbox
             coco_bbox = [bounds[0], bounds[1], bounds[2] - bounds[0], bounds[3] - bounds[1]]
-            coco_bbox = [round(xy, 2) for xy in coco_bbox]
+            coco_bbox = [round(coords, 2) for coords in coco_bbox]
 
             key_annotation = {"id": key_idx,
                               "image_id": int(chip_id),
@@ -77,7 +81,12 @@ def format_coco(set_: Dict, chip_width: int, chip_height: int):
 
 
 def move_coco_val_images(val_chips_list, path_train_folder):
-    """Move val chip images to val folder, applies train/val split on images"""
+    """Move validation chip images to val folder (applies train/val split on images)
+
+    Args:
+        val_chips_list: List of validation image key names
+        path_train_folder: Filepath to the training COCO image chip "train" folder
+    """
     out_folder = path_train_folder.parent / 'val2016'
     Path(out_folder).mkdir(parents=True, exist_ok=True)
     for chip in val_chips_list:
@@ -86,16 +95,15 @@ def move_coco_val_images(val_chips_list, path_train_folder):
 
 def coco_to_shapely(fp_coco_json: Union[Path, str],
                     categories: List[int]=None) -> Dict:
-    """
-    Transforms coco json annotations to shapely format.
+    """Transforms COCO annotations to shapely geometry format.
 
     Args:
         fp_coco_json: Input filepath coco json file.
         categories: Categories will filter to specific categories and images that contain at least one
         annotation of that category.
 
     Returns:
-        Dictionary of image key and shapely Multipolygon
+        Dictionary of image key and shapely Multipolygon.
     """
 
     data = utils.other.load_saved(fp_coco_json, file_format='json')
@@ -110,8 +118,9 @@ def coco_to_shapely(fp_coco_json: Union[Path, str],
     extracted_geometries = {}
     for image_id, file_name in zip(image_ids, file_names):
         annotations = [x for x in data['annotations'] if x['image_id'] == image_id]
-        # Filter to annotations of the selected category.
-        annotations = [x for x in annotations if x['category_id'] in categories]
+        if categories is not None:
+            annotations = [x for x in annotations if x['category_id'] in categories]
+
         segments = [segment['segmentation'][0] for segment in annotations]  # format [x,y,x1,y1,...]
 
         # Create shapely Multipolygons from COCO format polygons.
@@ -121,3 +130,19 @@ def coco_to_shapely(fp_coco_json: Union[Path, str],
     return extracted_geometries
 
 
+def plot_coco(in_json, chip_img_folder, start=0, end=2):
+    """Plot COCO annotations and image chips"""
+    extracted = utils.coco.coco_to_shapely(in_json)
+
+    for key in sorted(extracted.keys())[start:end]:
+        print(key)
+        plt.figure(figsize=(5, 5))
+        plt.axis('off')
+
+        img = np.asarray(pilimage.open(rf'{chip_img_folder}\{key}'))
+        plt.imshow(img, interpolation='none')
+
+        mp = extracted[key]
+        patches = [PolygonPatch(p, ec='r', fill=False, alpha=1, lw=0.7, zorder=1) for p in mp]
+        plt.gca().add_collection(PatchCollection(patches, match_original=True))
+        plt.show()

utils/coco.py

@@ -1,18 +1,18 @@
 # geo.py
 
 import warnings
-from typing import Union, Dict, Tuple, List
-import random
-import itertools
+from typing import Union, Dict
 
 import numpy as np
-import geopandas as gpd
 from geopandas import GeoDataFrame as GDF
 from pandas import DataFrame as DF
 import shapely
-from shapely.geometry import Polygon, MultiPolygon
+from shapely.geometry import Polygon
 import rasterio.crs
-from pathlib import Path
+import geopandas as gpd
+from tqdm import tqdm
+
+import utils.img
 
 
 def buffer_zero(in_geo: Union[GDF, Polygon]) -> Union[GDF, Polygon]:
@@ -54,6 +54,25 @@ def set_crs(df: GDF, epsg_code: Union[int, str]) -> GDF:
     return df
 
 
+def explode_mp(df: GDF) -> GDF:
+    """Explode all multi-polygon geometries in a geodataframe into individual polygon geometries.
+
+    Adds exploded polygons as rows at the end of the geodataframe and resets its index.
+    """
+    outdf = df[df.geom_type == 'Polygon']
+
+    df_mp = df[df.geom_type == 'MultiPolygon']
+    for idx, row in df_mp.iterrows():
+        df_temp = gpd.GeoDataFrame(columns=df_mp.columns)
+        df_temp = df_temp.append([row] * len(row.geometry), ignore_index=True)
+        for i in range(len(row.geometry)):
+            df_temp.loc[i, 'geometry'] = row.geometry[i]
+        outdf = outdf.append(df_temp, ignore_index=True)
+
+    outdf.reset_index(drop=True, inplace=True)
+    return outdf
+
+
 def clip(df: GDF,
          clip_poly: Polygon,
          explode_mp_: bool = False,
@@ -102,7 +121,7 @@ def reclassify_col(df: Union[GDF, DF],
                    ) -> Union[GDF, DF]:
     """Reclassify class label and class ids in a dataframe column.
 
-    # TODO: Make more efficient!
+    # TODO: Simplify & make more efficient!
     Args:
         df: input geodataframe.
         rcl_scheme: Reclassification scheme, e.g. {'springcereal': [1,2,3], 'wintercereal': [10,11]}
@@ -126,6 +145,7 @@ def reclassify_col(df: Union[GDF, DF],
 
     df[f'rcl_{col_classlabels}'] = df[col_classids].copy().map(rcl_dict)  # map name first, id second!
     df[f'rcl_{col_classids}'] = df[col_classids].map(rcl_dict_id)
+
     return df
 
 
@@ -250,4 +270,54 @@ def _invert_y_axis(poly: Polygon=ingeo, reference_height=reference_height):
         return _invert_y_axis(poly=ingeo, reference_height=reference_height)
     elif isinstance(ingeo, GDF):
         ingeo.geometry = ingeo.geometry.apply(lambda _p: _invert_y_axis(poly=_p, reference_height=reference_height))
-        return ingeo
+        return ingeo
+
+
+def cut_chip_geometries(vector_df, raster_width, raster_height, raster_transform, chip_width=128, chip_height=128,):
+    """Workflow to cut a vector geodataframe to chip geometries.
+
+    Filters small polygons and skips empty chips.
+
+    Args:
+        vector_df: Geodataframe containing the geometries to be cut to chip geometries.
+        raster_width: rasterio meta['width']
+        raster_height: rasterio meta['height']
+        raster_transform: rasterio meta['transform']
+        chip_width: Desired pixel width.
+        chip_height: Desired pixel height.
+
+    Returns: Dictionary containing the final chip_df, chip_window, chip_transform, chip_poly objects.
+    """
+
+    generator_window_bounds = utils.img.get_chip_windows(raster_width=raster_width,
+                                                         raster_height=raster_height,
+                                                         raster_transform=raster_transform,
+                                                         chip_width=chip_width,
+                                                         chip_height=chip_height,
+                                                         skip_partial_chips=True)
+
+    all_chip_dfs = {}
+    for i, (chip_window, chip_transform, chip_poly) in enumerate(tqdm(generator_window_bounds)):
+
+        # # Clip geometry to chip
+        chip_df = vector_df.pipe(utils.geo.clip, clip_poly=chip_poly, keep_biggest_poly_=True)
+        if not all(chip_df.geometry.is_empty):
+            chip_df.geometry = chip_df.simplify(1, preserve_topology=True)
+        else:
+            continue
+        # Drop small geometries
+        chip_df = chip_df[chip_df.geometry.area * (10 * 10) > 5000]  #5000 sqm in UTM
+        # Transform to chip pixelcoordinates and invert y-axis for COCO format.
+        if not all(chip_df.geometry.is_empty):
+            chip_df = chip_df.pipe(utils.geo.to_pixelcoords, reference_bounds=chip_poly.bounds, scale=True,
+                                   ncols=chip_width, nrows=chip_height)
+            chip_df = chip_df.pipe(invert_y_axis, reference_height=chip_height)
+        else:
+            continue
+
+        chip_name = f'COCO_train2016_000000{100000+i}'  # _{clip_minX}_{clip_minY}_{clip_maxX}_{clip_maxY}'
+        all_chip_dfs[chip_name] = {'chip_df': chip_df,
+                                   'chip_window': chip_window,
+                                   'chip_transform': chip_transform,
+                                   'chip_poly': chip_poly}
+    return all_chip_dfs

utils/geo.py

@@ -1,6 +1,6 @@
 # img.py
 
-from typing import Tuple, Generator, Union
+from typing import Tuple, Generator
 
 import rasterio.windows
 from rasterio.windows import Window
@@ -9,7 +9,6 @@
 import warnings
 from pathlib import Path
 
-import utils.geo
 import itertools
 import numpy as np
 import rasterio
@@ -63,56 +62,6 @@ def get_chip_windows(raster_width: int,
         yield (chip_window, chip_transform, chip_poly)
 
 
-def cut_chip_geometries(vector_df, raster_width, raster_height, raster_transform, chip_width=128, chip_height=128,):
-    """Workflow to cut a vector geodataframe to chip geometries.
-
-    Filters small polygons and skips empty chips.
-
-    Args:
-        vector_df: Geodataframe containing the geometries to be cut to chip geometries.
-        raster_width: rasterio meta['width']
-        raster_height: rasterio meta['height']
-        raster_transform: rasterio meta['transform']
-        chip_width: Desired pixel width.
-        chip_height: Desired pixel height.
-
-    Returns: Dictionary containing the final chip_df, chip_window, chip_transform, chip_poly objects.
-    """
-
-    generator_window_bounds = get_chip_windows(raster_width=raster_width,
-                                               raster_height=raster_height,
-                                               raster_transform=raster_transform,
-                                               chip_width=chip_width,
-                                               chip_height=chip_height,
-                                               skip_partial_chips=True)
-
-    all_chip_dfs = {}
-    for i, (chip_window, chip_transform, chip_poly) in enumerate(tqdm(generator_window_bounds)):
-
-        # # Clip geometry to chip
-        chip_df = vector_df.pipe(utils.geo.clip, clip_poly=chip_poly, keep_biggest_poly_=True)
-        if not all(chip_df.geometry.is_empty):
-            chip_df.geometry = chip_df.simplify(1, preserve_topology=True)
-        else:
-            continue
-        # Drop small geometries
-        chip_df = chip_df[chip_df.geometry.area * (10 * 10) > 5000]  #5000 sqm in UTM
-        # Transform to chip pixelcoordinates and invert y-axis for COCO format.
-        if not all(chip_df.geometry.is_empty):
-            chip_df = chip_df.pipe(utils.geo.to_pixelcoords, reference_bounds=chip_poly.bounds, scale=True,
-                                   ncols=chip_width, nrows=chip_height)
-            chip_df = chip_df.pipe(utils.geo.invert_y_axis, reference_height=chip_height)
-        else:
-            continue
-
-        chip_name = f'COCO_train2016_000000{100000+i}'  # _{clip_minX}_{clip_minY}_{clip_maxX}_{clip_maxY}'
-        all_chip_dfs[chip_name] = {'chip_df': chip_df,
-                                   'chip_window': chip_window,
-                                   'chip_transform': chip_transform,
-                                   'chip_poly': chip_poly}
-    return all_chip_dfs
-
-
 def cut_chip_images(img_path, chip_windows, out_folder, bands=[3, 2, 1]):
     """Cuts image chips and exports them
 
@@ -128,7 +77,7 @@ def cut_chip_images(img_path, chip_windows, out_folder, bands=[3, 2, 1]):
     src = rasterio.open(img_path)
 
     all_chip_stats = {}
-    for i, chip_window in enumerate(chip_windows):
+    for i, chip_window in enumerate(tqdm(chip_windows)):
         img_array = np.dstack(list(src.read(bands, window=chip_window)))
         img_array = exposure.rescale_intensity(img_array, in_range=(0, 2200))  # Sentinel2 range.
         with warnings.catch_warnings():

utils/img.py

@@ -1,6 +1,6 @@
 # other.py
 
-from typing import Union, Any, Callable, Tuple, Dict, Iterable, List
+from typing import Union, Any, Callable, Dict
 from pathlib import Path
 import pickle
 
@@ -40,7 +40,7 @@ def load_or_new_save(path: Path,
     """Write data to new pickle/json file or load pickle/json if that file already exists.
 
     Example:
-        df = cgeo.other.load_or_new_save(path=Path('output\preprocessed_marker_small.pkl'),
+        df = utils.other.load_or_new_save(path=Path('output\preprocessed_marker_small.pkl'),
                                          default_data=preprocess_vector,
                                          callable_args={'inpath': fp_fields, 'meta': meta})
     Args: