Aislar objetos de segmentación

Tras realizar la Tarea Segmento, a veces es conveniente extraer los objetos aislados de los resultados de la inferencia. Esta guía proporciona una receta genérica sobre cómo lograrlo utilizando elModo Predecir de Ultralytics .

Recorrido por las recetas

1. Comienza con las importaciones necesarias

   from pathlib import Path
   
   import cv2
   import numpy as np
   from ultralytics import YOLO
   

   Ultralytics Instala

   Consulta la sección Instalaciónrápida de Ultralytics para obtener una guía rápida sobre la instalación de las bibliotecas necesarias.

   ---

2. Carga un modelo y ejecútalo predict() en una fuente.

   from ultralytics import YOLO
   
   # Load a model
   model = YOLO('yolov8n-seg.pt')
   
   # Run inference
   results = model.predict()
   

   ¿No hay argumentos de predicción?

   Sin especificar una fuente, se utilizarán las imágenes de ejemplo de la biblioteca:

   'ultralytics/assets/bus.jpg'
   'ultralytics/assets/zidane.jpg'
   

   Esto es útil para realizar pruebas rápidas con el predict() método.

   Para obtener información adicional sobre los Modelos de Segmentación, visita la página Segmento Tarea página. Para saber más sobre predict() ver Modo Predecir de la Documentación.

   ---

3. Ahora itera sobre los resultados y los contornos. Para los flujos de trabajo que quieren guardar una imagen en un archivo, la imagen de origen base-name y la detección class-label se recuperan para su uso posterior (opcional).

   # (2) Iterate detection results (helpful for multiple images)
   for r in res:
       img = np.copy(r.orig_img)
       img_name = Path(r.path).stem # source image base-name
   
       # Iterate each object contour (multiple detections)
       for ci,c in enumerate(r):
           # (1) Get detection class name
           label = c.names[c.boxes.cls.tolist().pop()]
   

   1. Para saber más sobre cómo trabajar con los resultados de la detección, consulta la Sección de recuadros para el modo Predecir.
   2. Para saber más sobre predict() resultados ver Trabajar con los resultados del modo Predecir
   Bucle For

   Una sola imagen sólo iterará el primer bucle una vez. Una sola imagen con una sola detección iterará cada bucle sólo una vez.

   ---

4. Empieza por generar una máscara binaria a partir de la imagen de origen y luego dibuja un contorno relleno sobre la máscara. Esto permitirá aislar el objeto de las demás partes de la imagen. Un ejemplo de bus.jpg para uno de los detectados person se muestra a la derecha.

   

   # Create binary mask
   b_mask = np.zeros(img.shape[:2], np.uint8)
   
   # (1) Extract contour result
   contour = c.masks.xy.pop()
   # (2) Changing the type
   contour = contour.astype(np.int32)
   # (3) Reshaping
   contour = contour.reshape(-1, 1, 2)
   
   
   # Draw contour onto mask
   _ = cv2.drawContours(b_mask,
                       [contour],
                       -1,
                       (255, 255, 255),
                       cv2.FILLED)
   

   1. Para más información sobre c.masks.xy consulta Sección Máscaras del modo Predecir.

   2. Aquí los valores se introducen en np.int32 para compatibilidad con drawContours() de OpenCV.

   3. El OpenCV drawContours() espera que los contornos tengan una forma de [N, 1, 2] amplía la sección de abajo para más detalles.

   Ampliar para comprender lo que ocurre al definir la contour variable.

   • c.masks.xy :: Proporciona las coordenadas de los puntos del contorno de la máscara en el formato (x, y). Para más detalles, consulta Sección Máscaras del modo Predecir.

   • .pop() :: Como masks.xy es una lista que contiene un solo elemento, este elemento se extrae utilizando la función pop() método.

   • .astype(np.int32) :: Utilizando masks.xy devolverá con un tipo de datos de float32pero esto no será compatible con el OpenCV drawContours() así que esto cambiará el tipo de datos a int32 para la compatibilidad.

   • .reshape(-1, 1, 2) :: Reformatea los datos en la forma requerida de [N, 1, 2] donde N es el número de puntos del contorno, con cada punto representado por una única entrada 1y la entrada se compone de 2 valores. En -1 denota que el número de valores a lo largo de esta dimensión es flexible.

   Ampliar para una explicación de la drawContours() configuración.

   • Encapsular la contour variable entre corchetes, [contour]generó eficazmente la máscara de contorno deseada durante las pruebas.

   • El valor -1 especificado para el drawContours() indica a la función que dibuje todos los contornos presentes en la imagen.

   • En tuple (255, 255, 255) representa el color blanco, que es el color deseado para dibujar el contorno en esta máscara binaria.

   • La adición de cv2.FILLED coloreará todos los píxeles encerrados por el límite del contorno de la misma manera, en este caso, todos los píxeles encerrados serán blancos.

   • Consulta Documentación de OpenCV en drawContours() para más información.

   ---

5. A continuación hay 2 opciones sobre cómo avanzar con la imagen a partir de este punto y una opción posterior para cada una.

   Opciones de aislamiento de objetos

   Píxeles de fondo negroPíxeles de fondo transparente

   # Create 3-channel mask
   mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
   
   # Isolate object with binary mask
   isolated = cv2.bitwise_and(mask3ch, img)
   

   ¿Cómo funciona?

   • En primer lugar, la máscara binaria se convierte de una imagen de un solo canal a una imagen de tres canales. Esta conversión es necesaria para el paso posterior en el que se combinan la máscara y la imagen original. Ambas imágenes deben tener el mismo número de canales para que sean compatibles con la operación de mezcla.

   • La imagen original y la máscara binaria de tres canales se fusionan utilizando la función OpenCV bitwise_and(). Esta operación conserva sólo valores de píxel mayores que cero (> 0) de ambas imágenes. Como los píxeles de la máscara son mayores que cero (> 0) sólo dentro de la región del contorno, los píxeles que quedan de la imagen original son los que se solapan con el contorno.

   Aislar con píxeles negros: Subopciones

   Imagen a tamaño completo

   No se requieren pasos adicionales si mantienes la imagen a tamaño completo.

   

   Objeto recortado Imagen

   Pasos adicionales necesarios para recortar la imagen de modo que sólo incluya la región del objeto.

   

   # (1) Bounding box coordinates
   x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
   # Crop image to object region
   iso_crop = isolated[y1:y2, x1:x2]
   

   1. Para más información sobre los resultados de los recuadros delimitadores, consulta la sección Recuadros del modo Predecir

   ¿Qué hace este código?

   • En c.boxes.xyxy.cpu().numpy() recupera los recuadros delimitadores como una matriz NumPy en la función xyxy formato, donde xmin, ymin, xmaxy ymax representan las coordenadas del rectángulo de la caja delimitadora. Consulta Sección Cajas del Modo Predecir para más detalles.

   • En squeeze() elimina cualquier dimensión innecesaria de la matriz NumPy, garantizando que tenga la forma esperada.

   • Convertir los valores de las coordenadas mediante .astype(np.int32) cambia el tipo de datos de las coordenadas de la casilla de float32 a int32haciéndolas compatibles para el recorte de imágenes mediante cortes índice.

   • Por último, la región del cuadro delimitador se recorta de la imagen mediante el corte de índice. Los límites se definen mediante la función [ymin:ymax, xmin:xmax] coordenadas del cuadro delimitador de la detección.

   # Isolate object with transparent background (when saved as PNG)
   isolated = np.dstack([img, b_mask])
   

   ¿Cómo funciona?

   • Utilizar NumPy dstack() (apilamiento de matrices a lo largo del eje de profundidad) junto con la máscara binaria generada, creará una imagen con cuatro canales. Esto permite que todos los píxeles fuera del contorno del objeto sean transparentes al guardarlos como un PNG archivo.

   Aislar con píxeles transparentes: Subopciones

   Imagen a tamaño completo

   No se requieren pasos adicionales si mantienes la imagen a tamaño completo.

   

   Objeto recortado Imagen

   Pasos adicionales necesarios para recortar la imagen de modo que sólo incluya la región del objeto.

   

   # (1) Bounding box coordinates
   x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
   # Crop image to object region
   iso_crop = isolated[y1:y2, x1:x2]
   

   1. Para más información sobre los resultados de los recuadros delimitadores, consulta la sección Recuadros del modo Predecir

   ¿Qué hace este código?

   • Al utilizar c.boxes.xyxy.cpu().numpy()las cajas delimitadoras se devuelven como una matriz NumPy, utilizando la función xyxy formato de coordenadas de casilla, que corresponden a los puntos xmin, ymin, xmax, ymax para el cuadro delimitador (rectángulo), véase Sección Cajas del Modo Predecir para más información.

   • Añadir squeeze() garantiza que se elimine cualquier dimensión extraña de la matriz NumPy.

   • Convertir los valores de las coordenadas mediante .astype(np.int32) cambia el tipo de datos de las coordenadas de la casilla de float32 a int32 que será compatible al recortar la imagen mediante cortes índice.

   • Por último, la región de la imagen para el cuadro delimitador se recorta mediante el corte de índice, en el que los límites se establecen mediante la función [ymin:ymax, xmin:xmax] coordenadas del cuadro delimitador de la detección.

   ¿Y si quiero el objeto recortado incluyendo el fondo?

   Se trata de una función integrada en la biblioteca Ultralytics . Consulta la página save_crop argumento para Argumentos de Inferencia del Modo Predecir para más detalles.

   ---

6. Lo que se haga a continuación depende totalmente de ti como desarrollador. Se muestra un ejemplo básico de un posible paso siguiente (guardar la imagen en un archivo para utilizarla en el futuro).

   • NOTA: este paso es opcional y puedes omitirlo si no es necesario para tu caso de uso específico.
   Ejemplo Paso final

   # Save isolated object to file
   _ = cv2.imwrite(f'{img_name}_{label}-{ci}.png', iso_crop)
   

   • En este ejemplo, la img_name es el nombre base del archivo de imagen de origen, label es el nombre de la clase detectada, y ci es el índice de detección del objeto (en caso de varias instancias con el mismo nombre de clase).

Código de ejemplo completo

Aquí, todos los pasos de la sección anterior se combinan en un único bloque de código. Para un uso repetido, lo óptimo sería definir una función que realice algunos o todos los comandos contenidos en el bloque de código for-loops, pero eso es un ejercicio que se deja al lector.

from pathlib import Path

import cv2
import numpy as np
from ultralytics import YOLO

m = YOLO('yolov8n-seg.pt')#(4)!
res = m.predict()#(3)!

# iterate detection results (5)
for r in res:
    img = np.copy(r.orig_img)
    img_name = Path(r.path).stem

    # iterate each object contour (6)
    for ci,c in enumerate(r):
        label = c.names[c.boxes.cls.tolist().pop()]

        b_mask = np.zeros(img.shape[:2], np.uint8)

        # Create contour mask (1)
        contour = c.masks.xy.pop().astype(np.int32).reshape(-1, 1, 2)
        _ = cv2.drawContours(b_mask, [contour], -1, (255, 255, 255), cv2.FILLED)

        # Choose one:

        # OPTION-1: Isolate object with black background
        mask3ch = cv2.cvtColor(b_mask, cv2.COLOR_GRAY2BGR)
        isolated = cv2.bitwise_and(mask3ch, img)

        # OPTION-2: Isolate object with transparent background (when saved as PNG)
        isolated = np.dstack([img, b_mask])

        # OPTIONAL: detection crop (from either OPT1 or OPT2)
        x1, y1, x2, y2 = c.boxes.xyxy.cpu().numpy().squeeze().astype(np.int32)
        iso_crop = isolated[y1:y2, x1:x2]

        # TODO your actions go here (2)

1. La línea que rellena contour se combina aquí en una sola línea, donde antes se dividía en varias.
2. ¡Lo que vaya aquí depende de ti!
3. Consulta el Modo Predecir para obtener información adicional.
4. Para más información, consulta Tarea de segmento.
5. Más información sobre Trabajar con resultados
6. Más información sobre los Resultados de la Máscara de Segmentación

Comentarios