Skip to content

COCO Processing

create_dataset 

create_dataset(
    description,
    contributor,
    version=str(Version(major=0)),
    date_created=datetime.now(),
)

Instances and returns a new CocoDataset model with given kwargs.

Parameters:

Name Type Description Default
description str

Description of your COCO dataset

 required
contributor str

Main contributors of your COCO dataset, its images and its annotations

 required
version str

Initial SemVer version (defaults to 0.0.0)

 str(Version(major=0))
date_created datetime

Date when dataset was initially created, defaults to datetime.now()

 now()

Returns:

Type Description
CocoDataset

An instance of CocoDataset without Image- and Annotation objects
Source code in geococo\coco_manager.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def create_dataset(
    description: str,
    contributor: str,
    version: str = str(Version(major=0)),
    date_created: datetime = datetime.now(),
) -> CocoDataset:
    """Instances and returns a new CocoDataset model with given kwargs.

    :param description: Description of your COCO dataset
    :param contributor: Main contributors of your COCO dataset, its images and its
        annotations
    :param version: Initial SemVer version (defaults to 0.0.0)
    :param date_created: Date when dataset was initially created, defaults to
        datetime.now()
    :return: An instance of CocoDataset without Image- and Annotation objects
    """

    info = Info(
        version=version,
        description=description,
        contributor=contributor,
        date_created=date_created,
        year=date_created.year,
    )
    dataset = CocoDataset(info=info)
    return dataset

load_dataset 

load_dataset(json_path)

Dumps the contents of json_path as a string, interprets it as a CocoDataset model and returns it.

Parameters:

Name Type Description Default
json_path Path

path to the JSON file containing the json-encoded COCO dataset

 required

Returns:

Type Description
CocoDataset

An instance of CocoDataset with loaded Image- and Annotation objects from json_path
Source code in geococo\coco_manager.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
def load_dataset(json_path: pathlib.Path) -> CocoDataset:
    """Dumps the contents of json_path as a string, interprets it as a CocoDataset model
    and returns it.

    :param json_path: path to the JSON file containing the json-encoded COCO dataset
    :return: An instance of CocoDataset with loaded Image- and Annotation objects from
        json_path
    """

    with open(json_path, mode="r", encoding="utf-8") as json_fp:
        json_data = json_fp.read()
    dataset = CocoDataset.parse_raw(json_data)
    return dataset

save_dataset 

save_dataset(dataset, json_path)

JSON-encodes an instance of CocoDataset and saves it to json_path.

Parameters:

Name Type Description Default
dataset CocoDataset

An instance of CocoDataset

 required
json_path Path

where to save the JSON-encoded CocoDataset instance to

 required
Source code in geococo\coco_manager.py 
50
51
52
53
54
55
56
57
58
59
def save_dataset(dataset: CocoDataset, json_path: pathlib.Path) -> None:
    """JSON-encodes an instance of CocoDataset and saves it to json_path.

    :param dataset: An instance of CocoDataset
    :param json_path: where to save the JSON-encoded CocoDataset instance to
    """

    json_data = dataset.json()
    with open(json_path, mode="w", encoding="utf-8") as dst:
        dst.write(json_data)

append_dataset 

append_dataset(
    dataset,
    images_dir,
    src,
    window_bounds,
    labels,
    id_attribute=None,
    name_attribute=None,
    super_attribute=None,
)

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 attribute arguments are column names of the GeoDataFrame and are used to supply (super)category names, category ids, or both. If only names are given, ids are autogenerated (and vice versa). Lastly, each call to this method will increment the dataset version: patch if using the same image_path, minor if using a new image_path, and major if using a new output_dir.

Parameters:

Name Type Description Default
dataset CocoDataset

CocoDataset model to append images and annotations to

 required
images_dir Path

output directory for all label images

 required
src DatasetReader

open rasterio reader for input raster

 required
labels GeoDataFrame

GeoDataFrame containing labels and class_info ('category_id')

 required
window_bounds List[Tuple[int, int]]

a list of window_bounds to attempt to use ()

 required
id_attribute Optional[str]

Column containing category_id values

 None
name_attribute Optional[str]

Column containing category_name values

 None
super_attribute Optional[str]

Column containing supercategory values

 None

Returns:

Type Description
CocoDataset

CocoDataset with n appended Image, Category and Annotation instances
Source code in geococo\coco_processing.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def append_dataset(
    dataset: CocoDataset,
    images_dir: pathlib.Path,
    src: DatasetReader,
    window_bounds: List[Tuple[int, int]],
    labels: gpd.GeoDataFrame,
    id_attribute: Optional[str] = None,
    name_attribute: Optional[str] = None,
    super_attribute: Optional[str] = None,
) -> 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 
    attribute arguments are column names of the GeoDataFrame and are used to supply 
    (super)category names, category ids, or both. If only names are given, ids are 
    autogenerated (and vice versa). Lastly, each call to this method will increment 
    the dataset version: patch if using the same image_path, minor if using a new 
    image_path, and major if using a new output_dir.

    :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 window_bounds: a list of window_bounds to attempt to use ()
    :param id_attribute: Column containing category_id values
    :param name_attribute: Column containing category_name values
    :param super_attribute: Column containing supercategory values
    :return: CocoDataset with n appended Image, Category and Annotation instances
    """

    # checks presence, types and values of all required attributes
    labels = validate_labels(
        labels=labels,
        id_attribute=id_attribute,
        name_attribute=name_attribute,
        super_attribute=super_attribute,
    )

    # adding new Category instances from labels (if any)
    dataset.add_categories(
        category_ids=labels.get(id_attribute),
        category_names=labels.get(name_attribute),
        super_names=labels.get(super_attribute),
    )

    # updating labels with validated COCO keys (i.e. 'name', 'id', 'supercategory')
    labels = update_labels(
        labels=labels,
        categories=dataset.categories,
        id_attribute=id_attribute,
        name_attribute=name_attribute,
    )

    # Setting nodata and estimating window configuration
    parent_window = window_intersect(input_raster=src, input_vector=labels)
    date_created = get_date_created(raster_source=src)
    nodata_value = src.nodata if src.nodata else 0
    coco_profile = src.profile.copy()
    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 bump either minor or patch version
    dataset.add_source(source_path=pathlib.Path(src.name), date_captured=date_created)

    # bumps major version if images_dir was not used in this dataset before
    dataset.verify_used_dir(images_dir=images_dir)

    for child_window in tqdm(
        window_factory(parent_window=parent_window, schema=schema), total=n_windows
    ):
        # changing window to Shapely.geometry, used to check for intersecting labels
        window_geom = generate_window_polygon(datasource=src, window=child_window)
        intersect_mask = labels.intersects(window_geom)
        if not intersect_mask.any():
            continue

        # all_touched is needed because of np.ceil in WindowSource
        window_labels = labels[intersect_mask]
        window_image, window_transform = riomask(
            dataset=src, shapes=[window_geom], all_touched=True, crop=True
        )

        # padding is needed because rasterio clips any mask to the extent of the source
        window_shape = (src.count, child_window.width, child_window.height)
        if window_image.shape != window_shape:
            window_image = reshape_image(
                img_array=window_image, shape=window_shape, padding_value=nodata_value
            )

        # normalizing values to uint8 range (i.e COCO dtype)
        if window_image.dtype != np.uint8:
            window_image = cv2.normalize(
                window_image,
                None,
                alpha=0,
                beta=255,
                norm_type=cv2.NORM_MINMAX,
                dtype=cv2.CV_8U,  # type: ignore
            )

        # updating rasterio profile with window-specific values
        coco_profile.update(
            {
                "width": window_image.shape[1],
                "height": window_image.shape[2],
                "transform": window_transform,
            }
        )

        # saving window_image to disk (if it doesn't exist already)
        window_image_path = (
            images_dir / f"{dataset.next_source_id}_{child_window.col_off}_"
            f"{child_window.row_off}_{child_window.width}_{child_window.height}.jpg"
        )
        if not window_image_path.exists():
            with rasterio.open(window_image_path, "w", **coco_profile) as dst:
                dst.write(window_image)

        # Instancing and adding Image model to dataset (also bumps next_image_id)
        image_instance = Image(
            id=dataset.next_image_id,
            width=window_image.shape[1],
            height=window_image.shape[2],
            file_name=window_image_path,
            source_id=dataset.next_source_id,
            date_captured=date_created
        )

        # 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("id").iterrows():
                label_mask = mask_label(
                    input_raster=windowed_src, label=window_label.geometry
                )
                if not label_mask.any():
                    continue

                rle = cocomask.encode(np.asfortranarray(label_mask))
                bounding_box = cv2.boundingRect(label_mask.astype(np.uint8))
                area = np.sum(label_mask)
                iscrowd = 1 if isinstance(window_label.geometry, MultiPolygon) else 0

                annotation_instance = Annotation(
                    id=dataset.next_annotation_id,
                    image_id=dataset.next_image_id,
                    category_id=window_label["id"],
                    segmentation=rle,  # type: ignore
                    area=area,
                    bbox=bounding_box,
                    iscrowd=iscrowd,
                )

                dataset.add_annotation(annotation=annotation_instance)
        dataset.add_image(image=image_instance)
    return dataset