18 from typing
import Any, Optional, List, Dict
21 import xml.etree.ElementTree
as ET
24 from PIL.Image
import Image
as PILImage
25 from skimage
import measure
26 from shapely.geometry
import Polygon
30 from .shared
import getTag, getBoxes
31 from ....annotation
import CoretexSegmentationInstance, BBox
32 from ....dataset
import ImageDataset
35 ContourPoints = List[List[int]]
36 SegmentationPolygon = List[ContourPoints]
42 Represents the groundtruth object for contour matching
45 def __init__(self, object: ET.Element):
53 Represents the extracted contour from segmentation mask
56 def __init__(self, contours: ContourPoints, iou: Optional[float] =
None):
65 class InstanceExtractor:
67 def __init__(self, dataset: ImageDataset) ->
None:
68 self.__dataset = dataset
70 def createSubmasks(self, maskImage: PILImage) -> Dict[str, Any]:
72 Creates submasks for each segmentation mask
81 Dict[str, Any] -> Dictionary with submask image and color
84 width, height = maskImage.size
87 subMasks: Dict[str, Any] = {}
88 for x
in range(width):
89 for y
in range(height):
91 pixel = maskImage.getpixel((x, y))
94 if not isinstance(pixel, tuple):
95 raise TypeError(f
"Expected \"tuple\" recieved \"{type(pixel)}\"")
97 if not len(pixel) >= 3:
98 raise ValueError(f
"Expected pixel to has at least 3 channels (RGB).")
102 pixelStr = str(pixel)
103 subMask = subMasks.get(pixelStr)
105 subMasks[pixelStr] = Image.new(
'1', (width + 2, height + 2))
108 subMasks[pixelStr].putpixel((x + 1, y + 1), 1)
112 def reshapeContour(self, candidate: ContourCandidate) -> ContourPoints:
113 listOfPoints: ContourPoints = []
115 for segmentContour
in candidate.contours:
116 for value
in range(0, len(segmentContour) - 1, 2):
117 point = [segmentContour[value], segmentContour[value + 1]]
118 listOfPoints.append(point)
122 def calculateIoU(self, boxA: Dict[str, float], boxB: List[float]) -> float:
124 Calculates area of overlap for object boxes and contour boxes
128 boxA : Dict[str, float]
129 annotated object boxes
130 boxB : Dict[str, float]
131 extracted boxes from contour
135 Value of boxes area overlap
138 xmax = boxA[
'top_left_x'] + boxA[
'width']
139 ymax = boxA[
'top_left_y'] + boxA[
'height']
141 xA = max(boxA[
'top_left_x'], boxB[0])
142 yA = max(boxA[
'top_left_y'], boxB[1])
143 xB = min(xmax, boxB[2])
144 yB = min(ymax, boxB[3])
147 interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
150 boxAArea = (xmax - boxA[
'top_left_x'] + 1) * (ymax - boxA[
'top_left_y'] + 1)
151 boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
153 return interArea / float(boxAArea + boxBArea - interArea)
155 def boxesOverlap(self, bbox: Dict[str, Any], listOfPoints: ContourPoints) -> float:
157 Extracts boxes from contour and calculates IoU
161 bbox : Dict[str, Any]
162 annotated object boxes
163 listOfPoints : ContourPoints
168 float -> Calculated overlap of boxes area
171 xmin: Optional[float] =
None
172 ymin: Optional[float] =
None
173 xmax: Optional[float] =
None
174 ymax: Optional[float] =
None
176 for point
in listOfPoints:
180 if xmin
is None or xmax
is None:
184 if ymin
is None or ymax
is None:
189 xmin = min(xmin, currentX)
190 ymin = min(ymin, currentY)
191 xmax = max(xmax, currentX)
192 ymax = max(ymax, currentY)
194 if xmax
is not None and xmin
is not None and ymin
is not None and ymax
is not None:
195 contourBox = [xmin, ymin, xmax, ymax]
196 iou = self.calculateIoU(bbox, contourBox)
200 def matchContour(self, bbox: Dict[str, Any], objectCandidate: ObjectCandidate, contourCandidates: List[ContourCandidate]) -> ContourPoints:
202 Matches object and contour with max area of overlap
206 bbox : Dict[str, Any]
207 annotated object boxes
208 objectCandidate : ObjectCandidate
210 contourCandidates : List[ContourCandidate]
211 list of ContourCandidate objects
215 ContourPoints -> The corresponding contour which matches given object
221 for candidate
in contourCandidates:
222 if candidate.matched:
225 listOfPoints = self.reshapeContour(candidate)
226 candidate.iou = self.boxesOverlap(bbox, listOfPoints)
228 if maxIoU < candidate.iou:
229 maxIoU = candidate.iou
230 contourIndex = contourCandidates.index(candidate)
232 objectCandidate.matched =
True
233 contourCandidates[contourIndex].matched =
True
235 return contourCandidates[contourIndex].contours
237 def extractSubmaskContours(self, subMask: PILImage) -> ContourPoints:
239 Extracts contours from submask image
248 ContourPoints -> List of contours
251 subMaskArray = np.asarray(subMask)
252 contours = measure.find_contours(subMaskArray, 0.5)
254 segmentations: ContourPoints = []
255 for contour
in contours:
256 for i
in range(len(contour)):
257 row, col = contour[i]
258 contour[i] = (col - 1, row - 1)
261 poly = Polygon(contour)
264 if poly.geom_type ==
'MultiPolygon':
266 poly = poly.convex_hull
269 if poly.geom_type ==
'Polygon':
270 segmentation = np.array(poly.exterior.coords).ravel().tolist()
271 segmentations.append(segmentation)
275 def getSegmentationInstance(
277 objectCandidate: ObjectCandidate,
278 contourCandidates: Optional[List[ContourCandidate]]=
None
279 ) -> Optional[CoretexSegmentationInstance]:
281 label = getTag(objectCandidate.object,
"name")
285 clazz = self.__dataset.classByName(label)
289 bndbox = objectCandidate.object.find(
'bndbox')
293 boxes = getBoxes(bndbox)
297 polygon: ContourPoints = []
298 if contourCandidates
is not None:
299 polygon = self.matchContour(boxes, objectCandidate, contourCandidates)
301 bbox = BBox.decode(boxes)
303 if len(polygon) == 0:
304 polygon = [bbox.polygon]
306 return CoretexSegmentationInstance.create(clazz.classIds[0], bbox, polygon)
308 def __extractNonSegmentedInstances(self, objectCandidates: List[ObjectCandidate]) -> List[CoretexSegmentationInstance]:
309 coretexInstances: List[CoretexSegmentationInstance] = []
311 for objectCandidate
in objectCandidates:
312 instance = self.getSegmentationInstance(objectCandidate)
316 coretexInstances.append(instance)
318 return coretexInstances
320 def __extractSegmentedInstances(
322 segmentationPath: str,
324 objectCandidates: List[ObjectCandidate]
325 ) -> List[CoretexSegmentationInstance]:
327 imageFile = Image.open(os.path.join(segmentationPath, filename))
329 contourCandidates: List[ContourCandidate] = []
331 maskImage = imageFile.convert(
"RGB")
332 subMasks = self.createSubmasks(maskImage)
334 for color, subMask
in subMasks.items():
336 if color ==
'(224, 224, 192)':
339 annotation = self.extractSubmaskContours(subMask)
342 coretexInstances: List[CoretexSegmentationInstance] = []
344 for objectCandidate
in objectCandidates:
345 if objectCandidate.matched:
348 instance = self.getSegmentationInstance(objectCandidate, contourCandidates)
352 coretexInstances.append(instance)
354 return coretexInstances
356 def extractInstances(
360 segmentationPath: str
361 ) -> List[CoretexSegmentationInstance]:
363 Extracts polygons from segmentation masks and creates instances
368 file with annotations
369 objects : List[ET.Element]
373 objects = root.findall(
"object")
374 objectCandidates: List[ObjectCandidate] = [
ObjectCandidate(obj)
for obj
in objects]
376 segmented = getTag(root,
"segmented")
377 if segmented
is not None:
378 isSegmented = bool(int(segmented))
381 return self.__extractSegmentedInstances(segmentationPath, filename, objectCandidates)
383 return self.__extractNonSegmentedInstances(objectCandidates)
386 raise RuntimeError(
">> [Coretex] (segmented) XML tag missing")