By Olga Tsjernytska, Senior ingenieur machine learning

Native PyTorch- en TensorFlow-augmenters hebben een groot nadeel: ze kunnen niet tegelijkertijd een afbeelding en zijn segmentatiemasker, begrenzingsvak of keypoint-locaties vergroten. Er zijn dus twee opties: schrijf zelf functies of gebruik bibliotheken van derden. Ik heb beide geprobeerd, en de tweede optie is gewoon beter 🙂

Waarom Albumentaties?

 
Albumentaties was de eerste bibliotheek die ik heb geprobeerd, en ik ben ermee doorgegaan, omdat:

• Het is open source,
• Intuïtief,
• Snel,
• Heeft meer dan 60 verschillende augmentaties,
• Goed gedocumenteerd,
• En, wat het belangrijkste is, kan tegelijkertijd een afbeelding en zijn segmentatiemasker, begrenzingskader of hoofdpuntlocaties vergroten.

Er zijn nog twee vergelijkbare bibliotheken - beeld en Augmentor. Helaas kan ik geen vergelijking geven, omdat ik ze nog niet heb geprobeerd. Tot op dit moment was Albumentations net genoeg.
 

Korte zelfstudie

 
In deze korte zelfstudie laat ik zien hoe u afbeeldingen kunt vergroten voor segmentatie- en objectdetectietaken - eenvoudig met een paar regels code.

Als je deze tutorial wilt volgen:

1. Albumentaties installeren. Ik raad je echt aan om te controleren of je de nieuwste versie hebt, omdat oudere mogelijk bugs bevatten. Ik gebruik versie '1.0.0' en het werkt prima.
2. Download hieronder een testafbeelding met labels. Het is gewoon een willekeurige afbeelding van COCO-gegevensset. Ik heb het een beetje aangepast en opgeslagen in het formaat dat vereist is door Albumentations. Deze bibliotheek accepteert afbeeldingen als NumPy-arrays, segmentatiemaskers als NumPy-arrays en begrenzingsvakken als lijsten.

Downloaden

Laten we de afbeelding, het binaire pixelgewijze segmentatiemasker en een selectiekader laden. Het selectiekader wordt gedefinieerd als een lijst met 4 elementen - [x_min, y_min, breedte, hoogte].

import pickle
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches # load data
with open("image_data.pickle", "rb") as handle: image_data = pickle.load(handle) image = image_data["image"]
mask = image_data["mask"]
bbox = image_data["bbox_coco"] # visualize data
fig, ax = plt.subplots(1, 2, figsize=(12, 5))
ax[0].imshow(image)
ax[0].set_title("Image")
ax[1].imshow(image)
bbox_rect = patches.Rectangle( bbox[:2], bbox[2], bbox[3], linewidth=2, edgecolor="r", facecolor="none"
)
ax[1].add_patch(bbox_rect)
ax[1].imshow(mask, alpha=0.3, cmap="gray_r")
ax[1].set_title("Image + BBox + Mask")
plt.show()

Na het laden en visualiseren van de afbeelding zou je dit moeten krijgen:

Afbeelding. De uitvoer bij het uitvoeren van code voor afbeelding en de visualisatie van de labels.
Segmentatiemasker wordt gevisualiseerd als een transparant zwart-wit beeld (1 is zwart, 'paard').

Maskervergroting voor segmentatie. En nu kunnen we beginnen met Albumentations. Transformaties hier worden op dezelfde manier gedefinieerd als PyTorch en TensorFlow (Keras API):

• Definieer transformatie door verschillende augmentaties te combineren met een Compose-object.
• Elke augmentatie heeft argument 'p', de waarschijnlijkheid die moet worden toegepast, en daarnaast augmentatie-specifieke argumenten, zoals 'width' en 'height' voor RandomCrop.
• Gebruik gedefinieerde transformatie als een functie om de afbeelding en het masker te vergroten. Deze functie retourneert een woordenboek met de toetsen `image` en `mask`.

Hieronder staat de code voor het vergroten van de afbeelding (en het masker) met willekeurige 256 × 256 uitsnede (altijd) en horizontaal omdraaien (alleen in 50% gevallen).

import albumentations as A # define augmentation
transform = A.Compose([ A.RandomCrop(width=256, height=256, p=1), A.HorizontalFlip(p=0.5),
]) # augment and visualize images
fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6): transformed = transform(image=image, mask=mask)  ax[i // 3, i % 3].imshow(transformed["image"]) ax[i // 3, i % 3].imshow(transformed["mask"], alpha=0.3, cmap="gray_r")
plt.show()

Als resultaat zou je zoiets moeten krijgen. Uw vergrote afbeeldingen zullen anders zijn, omdat Albumentations willekeurige transformaties produceert. Raadpleeg voor een gedetailleerde zelfstudie over maskervergroting: originele documentatie.

Afbeelding. De uitvoer bij het uitvoeren van code voor gelijktijdige beeld- en maskervergroting.
Segmentatiemasker wordt gevisualiseerd als een transparant zwart-wit beeld (1 is zwart, 'paard')

Vergroting van begrenzingsvakken voor objectdetectie. Het is vergelijkbaar met augmentatie voor segmentatiemaskers, maar:

• Definieer bovendien `bbox_params`, waarbij u het formaat van het selectiekader en het argument voor de klassen van het selectiekader specificeert. 'coco' betekent begrenzingsvak in COCO-datasetformaat - [x_min, y_min, breedte, hoogte]. En het argument `bbox_classes` zal later worden gebruikt om klassen door te geven voor begrenzingsvakken.
• `transform` accepteert begrenzingsvakken als een lijst met lijsten. Bovendien vereist het begrenzingsvakklassen (als een lijst), zelfs als er een enkel begrenzingsvak in de afbeelding is.

Hieronder staat de code die RandomCrop en HorizonalFrip tegelijkertijd doet voor de afbeelding en het bijbehorende selectiekader.

# define augmentation
transform = A.Compose([ A.RandomCrop(width=256, height=256, p=1), A.HorizontalFlip(p=0.5),
], bbox_params=A.BboxParams(format='coco', label_fields=["bbox_classes"])) # augment and visualize images
bboxes = [bbox] #`transform` accepts bounding boxes as a list of lists.
bbox_classes = ["horse"] fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6): transformed = transform( image=image,  bboxes=bboxes,  bbox_classes=bbox_classes ) ax[i // 3, i % 3].imshow(transformed["image"]) trans_bbox = transformed["bboxes"][0] bbox_rect = patches.Rectangle( trans_bbox[:2], trans_bbox[2], trans_bbox[3], linewidth=2, edgecolor="r", facecolor="none", ) ax[i // 3, i % 3].add_patch(bbox_rect)
plt.show()

En hier zijn de resultaten. Voor het geval je een aantal specifieke begrenzingsvakvergrotingen nodig hebt – raadpleeg de originele documentatie.

Afbeelding. De uitvoer bij het uitvoeren van code voor gelijktijdige vergroting van afbeeldingen en begrenzingsvakken.

Gelijktijdige vergroting van meerdere doelen. Naast het gelijktijdig toestaan ​​van meerdere maskers vergroten or meerdere begrenzingsvakken, Albumentations heeft een functie om tegelijkertijd verschillende soorten labels te vergroten, bijvoorbeeld een masker en een selectiekader.

Als je een `transform` aanroept, geef hem dan gewoon alles wat je hebt:

# define augmentation
transform = A.Compose([ A.RandomCrop(width=256, height=256, p=1), A.HorizontalFlip(p=0.5),
], bbox_params=A.BboxParams(format='coco', label_fields=["bbox_classes"])) # augment and visualize images
bboxes = [bbox]
bbox_classes = ["horse"] fig, ax = plt.subplots(2, 3, figsize=(15, 10))
for i in range(6): transformed = transform( image=image,  mask=mask,  bboxes=bboxes,  bbox_classes=bbox_classes ) ax[i // 3, i % 3].imshow(transformed["image"]) trans_bbox = transformed["bboxes"][0] bbox_rect = patches.Rectangle( trans_bbox[:2], trans_bbox[2], trans_bbox[3], linewidth=2, edgecolor="r", facecolor="none", ) ax[i // 3, i % 3].add_patch(bbox_rect) ax[i // 3, i % 3].imshow(transformed["mask"], alpha=0.3, cmap="gray_r")
plt.show()

Uw resultaat ziet er uit zoals in de onderstaande afbeelding. En hier is meer gedetailleerde documentatie daarover.

Afbeelding. De uitvoer bij het uitvoeren van code voor een gelijktijdige afbeelding, segmentatiemasker en begrenzingsvakvergroting.
Segmentatiemasker wordt gevisualiseerd als een transparant zwart-wit beeld (1 is zwart, 'paard').

En meer. Albumentations heeft veel meer functies beschikbaar, zoals augmentatie voor sleutelpunten en Automatische vergroting. En het omvat ongeveer 60 verschillende augmentatietypes - letterlijk voor elke taak die u nodig heeft.

Compatibiliteit met PyTorch & TensorFlow

 
Hoogstwaarschijnlijk ga je Albumentations gebruiken als onderdeel van de PyTorch- of TensorFlow-trainingspijplijn, dus ik zal kort beschrijven hoe je dit moet doen.

PyTorch. Wanneer een aangepaste dataset maken, definieer Albumentations-transformatie in de functie `__init__` en roep deze aan in de functie `__getitem__`. PyTorch-modellen vereisen dat invoergegevens tensoren zijn, dus zorg ervoor dat u `ToTensorV2` als laatste stap toevoegt bij het definiëren van `transform` (een truc van een van de Albumentations-tutorials).

from torch.utils.data import Dataset
from albumentations.pytorch import ToTensorV2 class CustomDataset(Dataset): def __init__(self, images, masks): self.images = images # assume it's a list of numpy images self.masks = masks # assume it's a list of numpy masks self.transform = A.Compose([ A.RandomCrop(width=256, height=256, p=1), A.HorizontalFlip(p=0.5), ToTensorV2(), ]) def __len__(self): return len(self.images) def __getitem__(self, idx): image = self.images[idx] mask = self.masks[idx] transformed = self.transform(image=image, mask=mask) transformed_image = transformed["image"] transformed_mask = transformed["mask"] return transformed_image, transformed_mask

TensorFlow (Keras API) maakt het ook mogelijk om Aangepaste datasets, vergelijkbaar met PyTorch. Dus definieer Albumentations-transformatie in de functie `__init__` en roep het aan in de functie `__getitem__`. Vrij eenvoudig, niet?

from tensorflow import keras class CustomDataset(keras.utils.Sequence): def __init__(self, images, masks): self.images = images self.masks = masks self.batch_size = 1 self.img_size = (256, 256) self.transform = A.Compose([ A.RandomCrop(width=256, height=256, p=1), A.HorizontalFlip(p=0.5), ]) def __len__(self): return len(self.images) // self.batch_size def __getitem__(self, idx): """Returns a batch of samples""" i = idx * self.batch_size batch_images = self.images[i : i + self.batch_size] batch_masks = self.masks[i : i + self.batch_size] batch_images_stacked = np.zeros( (self.batch_size,) + self.img_size + (3,), dtype="uint8" ) batch_masks_stacked = np.zeros( (self.batch_size,) + self.img_size, dtype="float32" ) for i in range(len(batch_images)): transformed = self.transform( image=batch_images[i],  mask=batch_masks[i] ) batch_images_stacked[i] = transformed["image"] batch_masks_stacked[i] = transformed["mask"] return batch_images_stacked, batch_masks_stacked

Dat is het! Ik hoop dat deze tutorial je heeft aangemoedigd om Albumentations te proberen de volgende keer dat je werkt aan segmentatie, objectdetectie of keypoint-lokalisatietaak. Laat me weten of het gelukt is!

 
Bio: Olga Tsjernytska is een Senior Machine Learning Engineer bij een groot Oost-Europees outsourcingbedrijf; was betrokken bij verschillende datawetenschapsprojecten voor Amerikaanse, Europese en Aziatische topbedrijven; belangrijkste specialisatie en interesse is Deep Computer Vision.

ORIGINELE. Met toestemming opnieuw gepost.

Zie ook:

• De enige Jupyter Notebooks-extensie die u echt nodig hebt
• Een interactieve 3D-kaart maken en deze gemakkelijk met iedereen delen?
• Extractie van objecten in afbeeldingen en video's met behulp van 5 regels code