Comment tracer un masque de segmentation sur image en Python

Dans cet article, je te partage les fonctions que j’ai conçu pour tracer rapidement un masque de segmentation sur image en Python.

Une des tâches majeures de Computer Vision est la segmentation d’images.

La segmentation consiste à diviser une image en groupes de pixels appelés « segments ».

Pour faire de la segmentation, il faut déterminer les pixels appartenant à un segment.

Habituellement, un segment représente un objet dans une image.

En segmentant la totalité des pixels d’un objet, on peut facilement déterminer son emplacement et ses contours.

La précision nécessaire à la segmentation, en fait une tâche à la fois complexe et puissante dans la détection d’objets.

Ici, je te présenterai deux fonctions qui ont pour but de :

• tracer un seul masque de segmentation
• tracer plusieurs masques de segmentation en même temps

La dernière fonction est une extension de la première, par conséquent, même si tu possèdes plusieurs masques à tracer, je t’invite à suivre attentivement le développement de la première fonction.

Remarque : Dans les lignes qui suivent, je pars du principe que tu possède déjà les masques de segmentation d’une image. Si ce n’est pas le cas, tu peux consulter mon article sur le modèle de Deep Learning SAM sur ce lien pour apprendre à l’utiliser et obtenir tes propres masques de segmentation à partir d’une image.

La fonction pour tracer un seul masque de segmentation sur image

Pour tracer un masque de segmentation, nous allons utiliser les librairies open-cv et numpy.

On suppose que l’on possède deux tableau numpy :

• l’image avec une shape (w, h, 3)
• le masque avec une shape (w, h, 3)

Ici, le masque est un tableau de booléen. Il indique True si le pixel est dans le segment, False sinon.

Voila la fonction pour obtenir l’image segmenté :

import cv2
import numpy as np

def draw_mask(image, mask_generated) :
  masked_image = image.copy()

  masked_image = np.where(mask_generated.astype(int),
                          np.array([0,255,0], dtype='uint8'),
                          masked_image)

  masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

Dans la fonction, on crée tout d’abord une copie de l’image. Puis, on y applique le masque avec la fonction np.where().

On obtient alors une image possèdent un masque sur l’objet segmenté.

Malheureusement, le masque est opaque : on ne voit pas à travers. On ne peut alors pas distinguer l’objet segmenté.

Pour améliorer la visualisation, on superpose l’image de départ à l’image masquée (en modifiant leur opacité) avec le code cv2.addWeighted(image, 0.3, masked_image, 0.7, 0). On obtient alors une image qui préserve à la fois le masque et l’objet segmenté.

On peut alors afficher l’image ainsi :

segmented_image = draw_mask(image, mask)

cv2.imshow(segmented_image)

Avec la fonction draw_mask, on peut tracer un masque de segmentation sur notre image.

Voyons comment adapter cette fonction pour l’utiliser avec plusieurs masques.

La fonction pour tracer plusieurs masques de segmentation sur image

Les masques de segmentation sont souvent le résultat de modèle de Deep Learning. Dans ce contexte, deux types de résultats sont possibles :

• les masques sont contenus dans une liste de tableaux
• les masques sont contenus dans une liste de dictionnaires

Ici, nous allons adapter la fonction, développée dans la première partie, à ces deux cas particuliers.

À partir d’une liste de tableaux

Voici la fonction pour tracer plusieurs masques sur une image à partir d’une liste de tableaux :

def draw_masks_fromList(image, masks_generated, labels, colors) :
  masked_image = image.copy()
  for i in range(len(masks_generated)) :
    masked_image = np.where(np.repeat(masks_generated[i][:, :, np.newaxis], 3, axis=2),
                            np.asarray(colors[int(labels[i][-1])], dtype='uint8'),
                            masked_image)

    masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

On a ajouté à la fonction de base :

• une boucle for pour parcourir la liste de masques
• deux listes :
  • labels correspondant aux classes détectées
  • colors correspondant aux couleurs à utiliser pour chaque label

Ensuite, on peut afficher l’image :

COLORS = [(255,0,0),(255,255,0)]

LABELS = [[0], [0]] # 0 for bird, 1 for sheep

segmented_image = draw_mask(image, masks_list, LABELS, COLORS)

cv2.imshow(segmented_image)

On obtient une image segmentée à partir d’une liste de tableaux.

À partir d’une liste de dictionnaires

Voici la fonction pour tracer plusieurs masques sur une image à partir d’une liste de dictionnaires :

def draw_masks_fromDict(image, masks_generated) :
  masked_image = image.copy()
  for i in range(len(masks_generated)) :
    masked_image = np.where(np.repeat(masks_generated[i]['segmentation'].astype(int)[:, :, np.newaxis], 3, axis=2),
                            np.random.choice(range(256), size=3),
                            masked_image)

    masked_image = masked_image.astype(np.uint8)

  return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

La fonction est pensée pour des dictionnaires contenant les masques à la clé nommée 'segmentation' et, pour une utilisation sans classes préalables.

Ensuite, on peut afficher l’image :

segmented_image = draw_mask(image, masks_dict)

cv2.imshow(segmented_image)

On obtient une image segmentée à partir d’une liste de dictionnaires.

Comme tu peux le voir, tracer des masques de segmentation sur une image est relativement rapide. Néanmoins, il est possible que tu veuilles également tracer des masques de segmentation sur des vidéos.

Dans ce cas, tu peux accéder à ma formation gratuite Détection d’Objet Tout-Terrain. À l’intérieur tu apprendras à segmenter facilement des vidéos avec SAM.

Pour accéder à ta formation gratuite, clique-ici :

débloquer ma formation !