Merge pull request #5 from jaspersiebring/category_instances

Added mapping of arbitrary label attributes to COCO's `category_id`

README.md

@@ -158,8 +158,3 @@ session = fo.launch_app(coco_dataset, port=5151)
   <img src="https://github.com/jaspersiebring/GeoCOCO/assets/25051531/f8ab55da-b3cd-4beb-b082-7946e712ea5c" width="45%" height = 250/>
   <img src="https://github.com/jaspersiebring/GeoCOCO/assets/25051531/9a796a54-ffc2-49c3-95bc-59e5c0dd1d7c" width="45%" height = 250 />
 </p>
-
-
-# Planned features
-- [QGIS plugin](https://github.com/jaspersiebring/geococo-qgis-plugin).
-- Data visualization with `pycocotool`'s plotting functionality

geococo/coco_models.py

@@ -2,11 +2,12 @@
 import numpy as np
 import pathlib
 from datetime import datetime
-from typing import List, Optional
+from typing import List, Optional, Dict
 from typing_extensions import TypedDict
 
 from pydantic import BaseModel, ConfigDict, InstanceOf, model_validator
 from semver.version import Version
+from geococo.utils import assert_valid_categories
 
 
 class CocoDataset(BaseModel):
@@ -18,14 +19,21 @@ class CocoDataset(BaseModel):
     _next_image_id: int = 1
     _next_annotation_id: int = 1
     _next_source_id: int = 1
+    _category_mapper: Dict = {}
 
     @model_validator(mode="after")
     def _set_ids(self) -> CocoDataset:
         self._next_image_id = len(self.images) + 1
         self._next_annotation_id = len(self.annotations) + 1
         self._next_source_id = len(self.sources)
+        self._category_mapper = self._get_category_mapper()
         return self
 
+    def _get_category_mapper(self) -> Dict:
+        category_data = [(category.name, category.id) for category in self.categories]
+        category_mapper = dict(category_data) if category_data else {}
+        return category_mapper
+
     def add_annotation(self, annotation: Annotation) -> None:
         self.annotations.append(annotation)
         self._next_annotation_id += 1
@@ -47,6 +55,29 @@ def add_source(self, source_path: pathlib.Path) -> None:
 
         self._next_source_id = source.id
 
+    def add_categories(self, categories: np.ndarray) -> None:
+        # checking if categories are castable to str and under a certain size
+        categories = assert_valid_categories(categories=np.unique(categories))
+
+        # filtering existing categories
+        category_mask = np.isin(categories, list(self._category_mapper.keys()))
+        new_categories = categories[~category_mask]
+
+        # generating mapper from new categories
+        start = len(self._category_mapper.values()) + 1
+        end = start + new_categories.size
+        category_dict = dict(zip(new_categories, np.arange(start, end)))
+
+        # instance and append new Category objects to dataset
+        for category_name, category_id in category_dict.items():
+            category = Category(
+                id=category_id, name=str(category_name), supercategory="1"
+            )
+            self.categories.append(category)
+
+        # update existing category_mapper with new categories
+        self._category_mapper.update(category_dict)
+
     def bump_version(self, bump_method: str) -> None:
         bump_methods = ["patch", "minor", "major"]
         version = Version.parse(self.info.version)

geococo/coco_processing.py

@@ -27,29 +27,26 @@ def labels_to_dataset(
     src: DatasetReader,
     labels: gpd.GeoDataFrame,
     window_bounds: List[Tuple[int, int]],
+    category_attribute: str = "category_id",
 ) -> CocoDataset:
     """Move across a given geotiff, converting all intersecting labels to COCO
     annotations and appending them to a COCODataset model. This is done through
-    rasterio.Window objects, the bounds of which you can set with window_bounds
-    (also determines the size of the output images associated with the
-    Annotation instances). The degree of overlap between these windows is
-    determined by the dimensions of the given labels to maximize representation
-    in the resulting dataset.
-
-    The "iscrowd" attribute (see
-    https://cocodataset.org/#format-data)
-    is determined by whether    the respective labels are Polygon or
-    MultiPolygon instances. The "category_id" attribute,    which
-    represents class or category identifiers, is expected to be present
-    in the given labels    GeoDataFrame under the same name.
-
-    :param dataset: CocoDataset model to append images and annotations
-        to
+    rasterio.Window objects, the bounds of which you can set with window_bounds (also
+    determines the size of the output images associated with the Annotation instances).
+    The degree of overlap between these windows is determined by the dimensions of the
+    given labels to maximize representation in the resulting dataset.
+
+    The "iscrowd" attribute (see https://cocodataset.org/#format-data) is determined by
+    whether    the respective labels are Polygon or MultiPolygon instances. The
+    "category_id" attribute,    which represents class or category identifiers, is
+    expected to be present in the given labels    GeoDataFrame under the same name.
+
+    :param dataset: CocoDataset model to append images and annotations to
     :param images_dir: output directory for all label images
     :param src: open rasterio reader for input raster
-    :param labels: GeoDataFrame containing labels and class_info
-        ('category_id')
+    :param labels: GeoDataFrame containing labels and class_info ('category_id')
     :param window_bounds: a list of window_bounds to attempt to use ()
+    :param category_attribute: Column containing category_id values
     :return: The COCO dataset with appended Images and Annotations
     """
 
@@ -60,13 +57,16 @@ def labels_to_dataset(
     coco_profile.update({"dtype": np.uint8, "nodata": nodata_value, "driver": "JPEG"})
     schema = estimate_schema(gdf=labels, src=src, window_bounds=window_bounds)
     n_windows = generate_window_offsets(window=parent_window, schema=schema).shape[0]
-    
+
     # sets dataset.next_source_id and possibly bumps minor version
     dataset.add_source(source_path=pathlib.Path(src.name))
-    
+
     # bumps major version if images_dir has been used in this dataset before
     dataset.verify_new_output_dir(images_dir=images_dir)
-
+
+    # sets dataset._category_mapper
+    dataset.add_categories(categories=labels[category_attribute].unique())
+
     for child_window in tqdm(
         window_factory(parent_window=parent_window, schema=schema), total=n_windows
     ):
@@ -129,7 +129,9 @@ def labels_to_dataset(
 
         # Iteratively add Annotation models to dataset (also bumps next_annotation_id)
         with rasterio.open(window_image_path) as windowed_src:
-            for _, window_label in window_labels.sort_values("category_id").iterrows():
+            for _, window_label in window_labels.sort_values(
+                category_attribute
+            ).iterrows():
                 label_mask = mask_label(
                     input_raster=windowed_src, label=window_label.geometry
                 )
@@ -140,11 +142,13 @@ def labels_to_dataset(
                 bounding_box = cv2.boundingRect(label_mask.astype(np.uint8))
                 area = np.sum(label_mask)
                 iscrowd = 1 if isinstance(window_label.geometry, MultiPolygon) else 0
+                category_name = str(window_label[category_attribute])
+                category_id = dataset._category_mapper[category_name]
 
                 annotation_instance = Annotation(
                     id=dataset.next_annotation_id,
                     image_id=dataset.next_image_id,
-                    category_id=window_label["category_id"],
+                    category_id=category_id,
                     segmentation=rle,  # type: ignore
                     area=area,
                     bbox=bounding_box,

geococo/utils.py

@@ -13,14 +13,11 @@
 def mask_label(
     input_raster: DatasetReader, label: Union[Polygon, MultiPolygon]
 ) -> np.ndarray:
-    """Masks out an label from input_raster and flattens it to a 2D binary
-    array. If it doesn't overlap, the resulting mask will only consist of False
-    bools.
-
-    :param input_raster: open rasterio DatasetReader for the input
-        raster
-    :param label: Polygon object representing the area to be masked
-        (i.e. label)
+    """Masks out an label from input_raster and flattens it to a 2D binary array. If it
+    doesn't overlap, the resulting mask will only consist of False bools.
+
+    :param input_raster: open rasterio DatasetReader for the input raster
+    :param label: Polygon object representing the area to be masked (i.e. label)
     :return: A 2D binary array representing the label
     """
 
@@ -39,14 +36,12 @@ def mask_label(
 def window_intersect(
     input_raster: DatasetReader, input_vector: gpd.GeoDataFrame
 ) -> Window:
-    """Generates a Rasterio Window from the intersecting extents of the input
-    data. It also verifies if the input data share the same CRS and if they
-    physically overlap.
+    """Generates a Rasterio Window from the intersecting extents of the input data. It
+    also verifies if the input data share the same CRS and if they physically overlap.
 
     :param input_raster: rasterio dataset (i.e. input image)
     :param input_vector: geopandas geodataframe (i.e. input labels)
-    :return: rasterio window that represent the intersection between
-        input data extents
+    :return: rasterio window that represent the intersection between input data extents
     """
 
     if input_vector.crs != input_raster.crs:
@@ -73,13 +68,11 @@ def window_intersect(
 def reshape_image(
     img_array: np.ndarray, shape: Tuple[int, int, int], padding_value: int = 0
 ) -> np.ndarray:
-    """Reshapes 3D numpy array to match given 3D shape, done through slicing or
-    padding.
+    """Reshapes 3D numpy array to match given 3D shape, done through slicing or padding.
 
     :param img_array: the numpy array to be reshaped
     :param shape: the desired shape (bands, rows, cols)
-    :param padding_value: what value to pad img_array with (if too
-        small)
+    :param padding_value: what value to pad img_array with (if too small)
     :return: numpy array in desired shape
     """
 
@@ -98,14 +91,14 @@ def reshape_image(
 
 
 def generate_window_polygon(datasource: DatasetReader, window: Window) -> Polygon:
-    """Turns the spatial bounds of a given window to a shapely.Polygon object
-    in a given dataset's CRS.
+    """Turns the spatial bounds of a given window to a shapely.Polygon object in a given
+    dataset's CRS.
 
-    :param datasource: a rasterio DatasetReader object that provides the
-        affine transformation
+    :param datasource: a rasterio DatasetReader object that provides the affine
+        transformation
     :param window: bounds to represent as Polygon
-    :return: shapely Polygon representing the spatial bounds of a given
-        window in a given CRS
+    :return: shapely Polygon representing the spatial bounds of a given window in a
+        given CRS
     """
 
     window_transform = datasource.window_transform(window)
@@ -117,8 +110,7 @@ def generate_window_polygon(datasource: DatasetReader, window: Window) -> Polygo
 def generate_window_offsets(window: Window, schema: WindowSchema) -> np.ndarray:
     """Computes an array of window offsets bound by a given window.
 
-    :param window: the bounding window (i.e. offsets will be within its
-        bounds)
+    :param window: the bounding window (i.e. offsets will be within its bounds)
     :param schema: the parameters for the window generator
     :return: an array of window offsets within the bounds of window
     """
@@ -143,14 +135,14 @@ def generate_window_offsets(window: Window, schema: WindowSchema) -> np.ndarray:
 def window_factory(
     parent_window: Window, schema: WindowSchema, boundless: bool = True
 ) -> Generator[Window, None, None]:
-    """Generator that produces rasterio.Window objects in predetermined steps,
-    within the given Window.
+    """Generator that produces rasterio.Window objects in predetermined steps, within
+    the given Window.
 
-    :param parent_window: the window that provides the bounds for all
-        child_window objects
+    :param parent_window: the window that provides the bounds for all child_window
+        objects
     :param schema: the parameters that determine the window steps
-    :param boundless: whether the child_window should be clipped by the
-        parent_window or not
+    :param boundless: whether the child_window should be clipped by the parent_window or
+        not
     :yield: a rasterio.Window used for windowed reading/writing
     """
 
@@ -174,8 +166,7 @@ def estimate_average_bounds(
 ) -> Tuple[float, float]:
     """Estimates the average size of all features in a GeoDataFrame.
 
-    :param gdf: GeoDataFrame that contains all features (i.e.
-        shapely.Geometry objects)
+    :param gdf: GeoDataFrame that contains all features (i.e. shapely.Geometry objects)
     :param quantile: what quantile will represent the feature population
     :return: a tuple of floats representing average width and height
     """
@@ -195,19 +186,16 @@ def estimate_schema(
     quantile: float = 0.9,
     window_bounds: List[Tuple[int, int]] = [(256, 256), (512, 512)],
 ) -> WindowSchema:
-    """Attempts to find a schema that is able to represent the average
-    GeoDataFrame feature (i.e. sufficient overlap) but within the bounds given
-    by window_bounds.
-
-    :param gdf: GeoDataFrame that contains features that determine the
-        degree of overlap
-    :param src: The rasterio DataSource associated with the resulting
-        schema (i.e. bounds and pixelsizes)
+    """Attempts to find a schema that is able to represent the average GeoDataFrame
+    feature (i.e. sufficient overlap) but within the bounds given by window_bounds.
+
+    :param gdf: GeoDataFrame that contains features that determine the degree of overlap
+    :param src: The rasterio DataSource associated with the resulting schema (i.e.
+        bounds and pixelsizes)
     :param quantile: what quantile will represent the feature population
-    :param window_bounds: a list of possible limits for the window
-        generators
-    :return: (if found) a viable WindowSchema with sufficient overlap
-        within the window_bounds
+    :param window_bounds: a list of possible limits for the window generators
+    :return: (if found) a viable WindowSchema with sufficient overlap within the
+        window_bounds
     """
 
     # estimating the required overlap between windows for labels to be represented fully
@@ -241,3 +229,29 @@ def estimate_schema(
         ) from last_exception
 
     return schema
+
+
+def assert_valid_categories(
+    categories: np.ndarray, max_dtype: str = "<U50"
+) -> np.ndarray:
+    """Checks if all elements in categories array can be represented by strings of a
+    certain length (defaults to <U50)
+
+    :param categories: numpy array containing category values
+    :param max_dtype: numpy str dtype with char size
+    """
+
+    # checking if categories is castable to str (a prerequisite for class_names)
+    if not isinstance(categories, np.ndarray):
+        raise ValueError("Categories needs to be of type np.ndarray")
+
+    try:
+        str_categories = categories.astype(str)
+    except Exception as e:
+        raise ValueError("Category values need to be castable to str") from e
+
+    # checking if categories can be castable to str of a certain length (e.g. <U50)
+    if not np.can_cast(str_categories, max_dtype):
+        raise ValueError(f"Category values (str) have to fit in {max_dtype}")
+
+    return str_categories.astype(max_dtype)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "geococo"
-version = "0.2.1"
+version = "0.3.0"
 description = "Converts GIS annotations to Microsoft's Common Objects In Context (COCO) dataset format"
 authors = ["Jasper <[email protected]>"]
 readme = "README.md"