Datasets documentation
Segmentation sémantique
Segmentation sémantique
Les jeux de données de segmentation sémantique servent à entraîner un modèle à classifier chaque pixel d'une image. Ces jeux de données permettent une grande variété d'applications, telles que la suppression de l'arrière-plan d'images, la stylisation d'images ou la compréhension de scènes pour la conduite autonome. Ce guide vous montre comment appliquer des transformations à un jeu de données de segmentation d'images.
Avant de commencer, assurez-vous d'avoir installé des versions à jour de albumentations et cv2 :
pip install -U albumentations opencv-python
Albumentations est une bibliothèque Python pour effectuer l’augmentation de données en vision par ordinateur. Elle prend en charge diverses tâches de vision par ordinateur, telles que la classification d’images, la détection d’objets, la segmentation et l’estimation de points clés.
Ce guide utilise le jeu de données Scene Parsing pour segmenter et analyser une image en différentes régions associées à des catégories sémantiques, telles que le ciel, la route, la personne et le lit.
Chargez la partition train du jeu de données et examinez un exemple :
>>> from datasets import load_dataset
>>> dataset = load_dataset("scene_parse_150", split="train")
>>> index = 10
>>> dataset[index]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=683x512 at 0x7FB37B0EC810>,
'annotation': <PIL.PngImagePlugin.PngImageFile image mode=L size=683x512 at 0x7FB37B0EC9D0>,
'scene_category': 927}Le jeu de données comporte trois champs :
image: un objet d'image PIL.annotation: masque de segmentation de l'image.scene_category: l'étiquette ou la catégorie de scène de l'image (comme « cuisine » ou « bureau »).
Ensuite, consultez une image avec :
>>> dataset[index]["image"]
De même, vous pouvez consulter le masque de segmentation correspondant :
>>> dataset[index]["annotation"]
Nous pouvons également ajouter une palette de couleurs sur le masque de segmentation et le superposer à l’image originale pour visualiser le jeu de données :
Après avoir défini la palette de couleurs, vous devriez être prêt à visualiser certains superpositions.
>>> import matplotlib.pyplot as plt
>>> def visualize_seg_mask(image: np.ndarray, mask: np.ndarray):
... color_seg = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
... palette = np.array(create_ade20k_label_colormap())
... for label, color in enumerate(palette):
... color_seg[mask == label, :] = color
... color_seg = color_seg[..., ::-1] # convert to BGR
... img = np.array(image) * 0.5 + color_seg * 0.5 # plot the image with the segmentation map
... img = img.astype(np.uint8)
... plt.figure(figsize=(15, 10))
... plt.imshow(img)
... plt.axis("off")
... plt.show()
>>> visualize_seg_mask(
... np.array(dataset[index]["image"]),
... np.array(dataset[index]["annotation"])
... )
Appliquez maintenant quelques augmentations avec albumentations. Vous redimensionnez d’abord l’image et ajustez sa luminosité.
>>> import albumentations
>>> transform = albumentations.Compose(
... [
... albumentations.Resize(256, 256),
... albumentations.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.5),
... ]
... )Créez une fonction pour appliquer la transformation aux images :
>>> def transforms(examples):
... transformed_images, transformed_masks = [], []
...
... for image, seg_mask in zip(examples["image"], examples["annotation"]):
... image, seg_mask = np.array(image), np.array(seg_mask)
... transformed = transform(image=image, mask=seg_mask)
... transformed_images.append(transformed["image"])
... transformed_masks.append(transformed["mask"])
...
... examples["pixel_values"] = transformed_images
... examples["label"] = transformed_masks
... return examplesUtilisez la fonction set_transform() pour appliquer la transformation à la volée aux lots du jeu de données afin de réduire l’espace disque utilisé :
>>> dataset.set_transform(transforms)Vous pouvez vérifier que la transformation a fonctionné en accédant aux pixel_values et label d'un exemple :
>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])
>>> visualize_seg_mask(image, mask)
Dans ce guide, vous avez utilisé albumentations pour augmenter le jeu de données. Il est également possible d’utiliser torchvision pour appliquer des transformations similaires.
>>> from torchvision.transforms import Resize, ColorJitter, Compose
>>> transformation_chain = Compose([
... Resize((256, 256)),
... ColorJitter(brightness=0.25, contrast=0.25, saturation=0.25, hue=0.1)
... ])
>>> resize = Resize((256, 256))
>>> def train_transforms(example_batch):
... example_batch["pixel_values"] = [transformation_chain(x) for x in example_batch["image"]]
... example_batch["label"] = [resize(x) for x in example_batch["annotation"]]
... return example_batch
>>> dataset.set_transform(train_transforms)
>>> image = np.array(dataset[index]["pixel_values"])
>>> mask = np.array(dataset[index]["label"])
>>> visualize_seg_mask(image, mask)
Update on GitHubMaintenant que vous savez comment traiter un jeu de données pour la segmentation sémantique, apprenez comment entraîner un modèle de segmentation sémantique et l'utiliser pour l'inférence.