-
Notifications
You must be signed in to change notification settings - Fork 380
/
chip_classification_labels.py
180 lines (144 loc) · 5.9 KB
/
chip_classification_labels.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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
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
from typing import (TYPE_CHECKING, Any, Dict, Iterable, List, Optional,
Sequence, Tuple)
from dataclasses import dataclass
import numpy as np
from rastervision.core.box import Box
from rastervision.core.data.label import Labels
if TYPE_CHECKING:
from rastervision.core.data import (ClassConfig, CRSTransformer)
from shapely.geometry import Polygon
@dataclass
class ClassificationLabel:
class_id: int
scores: Optional[Sequence[float]] = None
def __iter__(self):
return iter((self.class_id, self.scores))
class ChipClassificationLabels(Labels):
"""Represents a spatial grid of cells associated with classes."""
def __init__(self,
cell_to_label: Optional[Dict[Box, Tuple[int, Optional[
Sequence[float]]]]] = None):
if cell_to_label is None:
cell_to_label = {}
self.cell_to_label = {
c: ClassificationLabel(*v)
for c, v in cell_to_label.items()
}
def __len__(self) -> int:
return len(self.cell_to_label)
def __eq__(self, other: 'ChipClassificationLabels') -> bool:
return (isinstance(other, ChipClassificationLabels)
and self.cell_to_label == other.cell_to_label)
def __add__(self, other: 'ChipClassificationLabels'
) -> 'ChipClassificationLabels':
result = ChipClassificationLabels()
result.extend(self)
result.extend(other)
return result
def __contains__(self, cell: Box) -> bool:
return cell in self.cell_to_label
def __getitem__(self, cell: Box) -> ClassificationLabel:
return self.cell_to_label[cell]
def __setitem__(self, window: Box,
value: Tuple[int, Optional[Sequence[float]]]):
class_id, scores = value
self.set_cell(window, class_id, scores=scores)
@classmethod
def from_predictions(cls, windows: Iterable['Box'],
predictions: Iterable[Any]) -> 'Labels':
"""Overrid to convert predictions to (class_id, scores) pairs."""
predictions = ((np.argmax(p), p) for p in predictions)
return super().from_predictions(windows, predictions)
@classmethod
def make_empty(cls) -> 'ChipClassificationLabels':
return ChipClassificationLabels()
def filter_by_aoi(self, aoi_polygons: Iterable['Polygon']):
result = ChipClassificationLabels()
for cell in self.cell_to_label:
cell_box = Box(*cell)
cell_poly = cell_box.to_shapely()
for aoi in aoi_polygons:
if cell_poly.within(aoi):
(class_id, scores) = self.cell_to_label[cell]
result.set_cell(cell_box, class_id, scores)
return result
def set_cell(self,
cell: Box,
class_id: int,
scores: Optional['np.ndarray'] = None) -> None:
"""Set cell and its class_id.
Args:
cell: (Box)
class_id: int
scores: 1d numpy array of probabilities for each class
"""
if scores is not None:
scores = list(map(lambda x: float(x), list(scores)))
class_id = int(class_id)
self.cell_to_label[cell] = ClassificationLabel(class_id, scores)
def get_cell_class_id(self, cell: Box) -> int:
"""Return class_id for a cell.
Args:
cell: (Box)
"""
result = self.cell_to_label.get(cell)
if result is not None:
return result.class_id
else:
return None
def get_cell_scores(self, cell: Box) -> Optional[Sequence[float]]:
"""Return scores for a cell.
Args:
cell: (Box)
"""
result = self.cell_to_label.get(cell)
if result is not None:
return result.scores
else:
return None
def get_singleton_labels(self, cell: Box):
"""Return Labels object representing a single cell.
Args:
cell: (Box)
"""
return ChipClassificationLabels({cell: self[cell]})
def get_cells(self) -> List[Box]:
"""Return list of all cells (list of Box)."""
return list(self.cell_to_label.keys())
def get_class_ids(self) -> List[int]:
"""Return list of class_ids for all cells."""
return [label.class_id for label in self.cell_to_label.values()]
def get_scores(self) -> List[Optional[Sequence[float]]]:
"""Return list of scores for all cells."""
return [label.scores for label in self.cell_to_label.values()]
def get_values(self) -> List[ClassificationLabel]:
"""Return list of class_ids and scores for all cells."""
return list(self.cell_to_label.values())
def extend(self, labels: 'ChipClassificationLabels') -> None:
"""Adds cells contained in labels.
Args:
labels: ChipClassificationLabels
"""
for cell in labels.get_cells():
self.set_cell(cell, *labels[cell])
def save(self,
uri: str,
class_config: 'ClassConfig',
crs_transformer: 'CRSTransformer',
bbox: Optional[Box] = None) -> None:
"""Save labels as a GeoJSON file.
Args:
uri (str): URI of the output file.
class_config (ClassConfig): ClassConfig to map class IDs to names.
crs_transformer (CRSTransformer): CRSTransformer to convert from
pixel-coords to map-coords before saving.
bbox (Optional[Box]): User-specified crop of the extent. Must be
provided if the corresponding RasterSource has bbox != extent.
"""
from rastervision.core.data import ChipClassificationGeoJSONStore
label_store = ChipClassificationGeoJSONStore(
uri=uri,
class_config=class_config,
crs_transformer=crs_transformer,
bbox=bbox)
label_store.save(self)