Segmentar el fondo de la imagen (extraer el cuerpo humano)

¡Este artículo utiliza principalmente rembg para extraer el cuerpo humano de la imagen y obtiene la máscara correspondiente y la imagen del cuerpo humano extraída! rembg es una buena herramienta de eliminación de fondo que utiliza la red UNet para entrenar y obtener un modelo de segmentación.

Configuración del entorno

python==3.8
#torch根据自己的[CUDA版本安装](https://pytorch.org/get-started/previous-versions/)
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install
			argparse
			opencv-python
			rembg
			pillow

código

Infer_mask.py
Este código incluye principalmente la función principal, la función de carga de imágenes, la función de adquisición de máscara y la función de extracción del cuerpo humano.

import cv2
import numpy as np
import os 
import io
import argparse
import torch
from torchvision import transforms
from PIL import Image
from rembg.bg import remove

def load_img(img_file):

    img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    if not img_file.endswith("png"):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)

    return img

def get_image_mask(img_file):

    mask_to_origin_tensor = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.0, ), (1.0, ))
    ])
    img_ori = load_img(img_file)
    with torch.no_grad():
        buf = io.BytesIO()
        Image.fromarray(img_ori).save(buf, format='png')
        img_pil = Image.open(
            io.BytesIO(remove(buf.getvalue()))).convert("RGBA")
    img_mask = torch.tensor(1.0) - (mask_to_origin_tensor(img_pil.split()[-1]) <
                                    torch.tensor(0.5)).float()

    return img_mask

def apply_mask(image, mask):
    """Apply the given mask to the image.
    """
    for c in range(3):
        image[:, :, c] = np.where(mask == 0,
                                  255,
                                  image[:, :, c])
    return image

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--img_path", type=str, default="./examples/5a6a25963db2f667441d5076972c207c.png")
    parser.add_argument("--save_path", type=str, default="./results/")
    parser.add_argument("--save_seg", action="store_true")
    args = parser.parse_args()

    img = cv2.imread(args.img_path, cv2.IMREAD_UNCHANGED)
    base=os.path.basename(args.img_path)
    name=os.path.splitext(base)[0]
    mask=get_image_mask(args.img_path)

    if args.save_seg:
        cv2.imwrite(args.save_path+name+'_image.png',apply_mask(img,np.squeeze(mask.numpy(),0)))
    cv2.imwrite(args.save_path+name+'_mask.png',np.squeeze(mask.numpy(),0)*255)
    print(args.save_path+name+'_mask.png')```

correr

Agregar save_seg guardará la imagen del cuerpo humano segmentada.

##################run single################
#To get a mask, result will be save at './results'
python Infer_mask.py --img_path ./00025_0052.png --save_path ./ --save_seg

############Run batch mask#############
image_path='/data1/supersmpl/dy/ReCurrence/A_Tools/Mask' #要处理图片的文件夹
save_path='./'  #结尾记得加反斜杠
for file_a in ${image_path}/*g; 
    do           #所有图像
        echo 'Process' `basename $file_a`
        python Infer_mask.py --img_path $file_a --save_path $save_path #--save_seg
    done

Aquí hay un ejemplo dividido:

pregunta

Si su programa tarda mucho en ejecutarse o se congela directamente, ¡generalmente es un problema con la conexión del calcetín! ¡Puede descargar manualmente el modelo previamente entrenado y colocarlo en la carpeta .u2net debajo de la carpeta de inicio! Enlace de descarga

charlar

Tengo algunas sugerencias sobre los estándares de publicación de artículos de CSDN. Cada vez que me recuerdan que "la calidad del artículo es baja, la recomendación se verá afectada", mi corazón late con fuerza. Después de todo, lo escribí yo mismo. Cuando recibí el mensaje, no tenía idea de por qué lo recibí. Más tarde, no tuve más remedio que consultar con el servicio de atención al cliente. El servicio de atención al cliente dijo que era porque había muy pocas palabras en este capítulo y me sugirieron que agregara algunas palabras. Sin embargo, principalmente compartí el código. Agregué comentarios de código para las necesidades necesarias. explicaciones de texto. No sé cómo cambiarlo, así que tengo una pequeña charla en la parte de atrás. Espero que CSDN siga siendo cauteloso al dar este consejo. Los creadores ponen sus esfuerzos en cada artículo. ¡Esto es realmente incómodo!

Referencia

Xiu, Y., Yang, J., Tzionas, D. y Black, MJ (2022). Icono: Humanos vestidos implícitamente obtenidos de normales. En Actas de la Conferencia IEEE/CVF sobre visión por computadora y reconocimiento de patrones (págs. 13296-13306).