18 from typing
import List, Tuple
19 from typing_extensions
import Self
20 from pathlib
import Path
21 from zipfile
import ZipFile
27 from PIL.Image
import Image
as PILImage
28 from numpy
import ndarray
32 from .base
import BaseImageDataset
33 from ...sample
import ImageSample, AnnotatedImageSampleData
34 from ...annotation
import CoretexSegmentationInstance, CoretexImageAnnotation, BBox
35 from ...._folder_manager
import folder_manager
38 ANNOTATION_NAME =
"annotations.json"
41 class AugmentedImageSample(ImageSample):
44 def path(self) -> Path:
48 Path -> path to new augmented samples directory
51 return folder_manager.temp /
"temp-augmented-ds" / str(self.id)
54 def createFromSample(cls, sample: ImageSample) -> Self:
56 Creates exact copy of sample from provided sample
70 for key, value
in sample.__dict__.items():
71 obj.__dict__[key] = copy.deepcopy(value)
80 locations: List[Tuple[int, int, int, int]]
84 if (x < loc[0] + loc[2]
and x + image.width > loc[0]
and
85 y < loc[1] + loc[3]
and y + image.height > loc[1]):
92 def generateSegmentedImage(image: np.ndarray, segmentationMask: np.ndarray) -> PILImage:
93 rgbaImage = Image.fromarray(image).convert(
"RGBA")
95 segmentedImage = np.asarray(rgbaImage) * segmentationMask
96 segmentedImage = Image.fromarray(segmentedImage)
98 alpha = segmentedImage.getchannel(
"A")
99 bbox = alpha.getbbox()
100 croppedImage = segmentedImage.crop(bbox)
102 if not isinstance(croppedImage, PILImage):
103 raise TypeError(f
"Expected \"PIL.Image.Image\" recieved \"{type(croppedImage)}\"")
109 segmentedImages: List[PILImage],
110 backgroundImage: np.ndarray,
113 ) -> Tuple[PILImage, List[Tuple[int, int]]]:
115 centroids: List[Tuple[int, int]] = []
116 locations: List[Tuple[int, int, int, int]] = []
118 background = Image.fromarray(backgroundImage)
120 for segmentedImage
in segmentedImages:
121 image = segmentedImage
123 rotatedImage = image.rotate(angle, expand =
True)
124 resizedImage = rotatedImage.resize((int(rotatedImage.width * scale), int(rotatedImage.height * scale)))
127 maxX = background.width - resizedImage.width
128 maxY = background.height - resizedImage.height
132 x = np.random.randint(0, maxX)
133 y = np.random.randint(0, maxY)
136 if not isOverlapping(x, y, resizedImage, locations):
139 background.paste(resizedImage, (x, y), resizedImage)
141 centerX = x + resizedImage.width // 2
142 centerY = y + resizedImage.height // 2
144 centroids.append((centerX, centerY))
147 locations.append((x, y, resizedImage.width, resizedImage.height))
149 return background, centroids
154 backgroundSampleData: AnnotatedImageSampleData,
157 ) -> Tuple[PILImage, List[CoretexSegmentationInstance]]:
159 segmentedImages: List[PILImage] = []
160 augmentedInstances: List[CoretexSegmentationInstance]= []
162 sampleData = sample.load()
163 if sampleData.annotation
is None:
164 raise RuntimeError(f
"CTX sample dataset sample id: {sample.id} image doesn't exist!")
166 for instance
in sampleData.annotation.instances:
167 foregroundMask = instance.extractBinaryMask(sampleData.image.shape[1], sampleData.image.shape[0])
168 segmentedImage = generateSegmentedImage(sampleData.image, foregroundMask)
170 segmentedImages.append(segmentedImage)
172 composedImage, centroids = composeImage(segmentedImages, backgroundSampleData.image, angle, scale)
174 for i, instance
in enumerate(sampleData.annotation.instances):
175 segmentationsFlatten = [sample
for sublist
in instance.segmentations
for sample
in sublist]
176 augmentedInstance = CoretexSegmentationInstance.create(instance.classId, BBox.fromPoly(segmentationsFlatten), instance.segmentations)
178 augmentedInstance.rotateSegmentations(angle)
179 augmentedInstance.centerSegmentations(centroids[i])
181 augmentedInstances.append(augmentedInstance)
183 return composedImage, augmentedInstances
188 backgroundSample: ImageSample,
191 ) -> Tuple[PILImage, CoretexImageAnnotation]:
193 backgroundSampleData = backgroundSample.load()
195 composedImage, augmentedInstances = processInstance(sample, backgroundSampleData, angle, scale)
196 annotation = CoretexImageAnnotation.create(sample.name, composedImage.width, composedImage.height, augmentedInstances)
198 return composedImage, annotation
203 augmentedSample: ImageSample,
204 augmentedImage: PILImage,
205 annotation: CoretexImageAnnotation
208 imagePath = tempPath / f
"{augmentedSample.name}.jpeg"
209 annotationPath = tempPath / ANNOTATION_NAME
210 zipPath = tempPath / f
"{augmentedSample.id}"
212 augmentedImage.save(imagePath)
214 with open(annotationPath,
'w')
as annotationfile:
215 jsonObject = json.dumps(annotation.encode())
216 annotationfile.write(jsonObject)
218 with ZipFile(zipPath.with_suffix(
".zip"), mode =
"w")
as archive:
219 archive.write(imagePath, f
"{augmentedSample.name}.jpeg")
220 archive.write(annotationPath, ANNOTATION_NAME)
222 imagePath.unlink(missing_ok =
True)
223 annotationPath.unlink(missing_ok =
True)
227 normalDataset: BaseImageDataset,
228 backgroundDataset: BaseImageDataset,
233 Modifies normalDataset by adding new augmented samples to it
237 normalDataset : BaseImageDataset
238 BaseImageDataset object
239 backgroundDataset : BaseImageDataset
240 BaseImageDataset object
242 angle of rotation in degrees
247 tempPath = folder_manager.createTempFolder(
"temp-augmented-ds")
248 augmentedSamples: List[AugmentedImageSample] = []
250 for i, background
in enumerate(backgroundDataset.samples):
253 for j, sample
in enumerate(normalDataset.samples):
255 augmentedImage, annotations = processSample(sample, background, angle, scale)
257 augmentedSample = AugmentedImageSample.createFromSample(sample)
258 augmentedSample.id = int(f
"{i}{j}{augmentedSample.id}")
260 storeFiles(tempPath, augmentedSample, augmentedImage, annotations)
262 augmentedSamples.append(augmentedSample)
264 for sample
in augmentedSamples:
265 normalDataset.samples.append(sample)