Creación de un conjunto de datos personalizado para FoundationPose y ejecución exitosa de un demo

Este artículo describe el proceso de generar un conjunto de datos personalizado para FoundationPose y ejecutar un demo con éxito, enfocándose en la adquisición de datos con una cámara Intel RealSense y la preparación de archivos necesarios.

  1. Preparación del entorno

Para configurar el entorno, se recomienda seguir las instrucciones específicas para WSL2 o utilizar Docker según las indicaciones del repositorio oficial.

  1. Análisis del conjunto de datos de FoundationPose

El conjunto de datos se compone de cinco elementos esenciales:

  • depth: Carpeta para almacenar mapas de profundidad capturados por la cámara.
  • masks: Carpeta que contiene una máscara de la primera imagen de la cámara RGB. El nombre del archivo debe coincidir con la primera imagen en las carpetas depth y rgb para evitar errores.
  • mesh: Carpeta con los archivos del modelo 3D, incluyendo el archivo OBJ, el archivo MTL y la textura en formato PNG.
  • rgb: Carpeta para las imágenes a color capturadas.
  • cam_K.txt: Archivo de texto con la matriz de parámetros intrínsecos de la cámara RGB.
  1. Creación de un conjunto de datos personalizado

3.1 Captura de imágenes RGB y de profundidad

Se utiliza una cámara Intel RealSense D435 para la adquisición de datos. El siguiente código en Python configura la cámara y guarda las imágenes con marcas de tiempo:

import pyrealsense2 as rs
import numpy as np
import cv2
import os
from datetime import datetime

def configurar_camara():
    cam_pipeline = rs.pipeline()
    config = rs.config()
    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
    config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
    perfil = cam_pipeline.start(config)
    sensor_depth = perfil.get_device().first_depth_sensor()
    escala_depth = sensor_depth.get_depth_scale()
    alineador = rs.align(rs.stream.color)
    return cam_pipeline, alineador, escala_depth

def capturar_datos_rgbd(dir_salida, num_imagenes=1000, intervalo=0.5):
    dir_rgb = os.path.join(dir_salida, "rgb")
    dir_depth = os.path.join(dir_salida, "depth")
    os.makedirs(dir_rgb, exist_ok=True)
    os.makedirs(dir_depth, exist_ok=True)
    
    cam_pipeline, alineador, escala_depth = configurar_camara()
    frames = cam_pipeline.wait_for_frames()
    frames_alineados = alineador.process(frames)
    frame_color = frames_alineados.get_color_frame()
    intrinsecos = frame_color.profile.as_video_stream_profile().intrinsics
    parametros = {
        "ancho": intrinsecos.width,
        "alto": intrinsecos.height,
        "fx": intrinsecos.fx,
        "fy": intrinsecos.fy,
        "ppx": intrinsecos.ppx,
        "ppy": intrinsecos.ppy,
        "escala_depth": escala_depth
    }
    
    with open(os.path.join(dir_salida, "parametros_camara.json"), "w") as archivo:
        import json
        json.dump(parametros, archivo, indent=4)
    
    contador = 0
    while contador < num_imagenes:
        frames = cam_pipeline.wait_for_frames()
        frames_alineados = alineador.process(frames)
        frame_depth = frames_alineados.get_depth_frame()
        frame_color = frames_alineados.get_color_frame()
        if not frame_depth or not frame_color:
            continue
        
        imagen_depth = np.asanyarray(frame_depth.get_data())
        imagen_color = np.asanyarray(frame_color.get_data())
        marca_temporal = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        
        cv2.imwrite(os.path.join(dir_rgb, f"img_{marca_temporal}.png"), imagen_color)
        cv2.imwrite(os.path.join(dir_depth, f"img_{marca_temporal}.png"), imagen_depth)
        
        contador += 1
        mapa_profundidad = cv2.applyColorMap(cv2.convertScaleAbs(imagen_depth, alpha=0.03), cv2.COLORMAP_JET)
        imagenes_combinadas = np.hstack((imagen_color, mapa_profundidad))
        cv2.imshow('Vista previa', imagenes_combinadas)
        if cv2.waitKey(int(intervalo * 1000)) & 0xFF == ord('q'):
            break
    
    cam_pipeline.stop()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    directorio = f"captura_realsense_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    capturar_datos_rgbd(directorio, num_imagenes=2000, intervalo=0.001)

3.2 Obtención de la matriz de parámetros intrínsecos

Los parámetros intrínsecos se guardan en un archivo JSON durante la captura. Se convierten a un archivo de texto compatible con FoundationPose usando este script:

import json

def generar_matriz_intrinsecos(archivo_json="parametros_camara.json", salida="cam_K.txt"):
    with open(archivo_json, 'r') as f:
        datos = json.load(f)
    
    fx = datos["fx"]
    fy = datos["fy"]
    ppx = datos["ppx"]
    ppy = datos["ppy"]
    
    matriz = [[fx, 0.0, ppx], [0.0, fy, ppy], [0.0, 0.0, 1.0]]
    with open(salida, 'w') as f:
        for fila in matriz:
            linea = " ".join([f"{val:.18e}" for val in fila])
            f.write(linea + "\n")

if __name__ == "__main__":
    generar_matriz_intrinsecos()

3.3 Creación de la máscara para la primera imagen

Se utiliza el framework Grounded-Segment-Anything para generar una máscara de segmentación. A continuación, un código adaptado que extrae la región de interés y la convierte en una imagen binaria:

import torch
import cv2
import numpy as np
from segment_anything import sam_model_registry, SamPredictor

def cargar_y_predecir_mascara(ruta_imagen, texto_prompt, checkpoint_sam, dispositivo="cuda"):
    # Carga del modelo SAM (ajustar según configuración de Grounded-Segment-Anything)
    modelo_sam = sam_model_registry["vit_b"](checkpoint=checkpoint_sam).to(dispositivo)
    predictor = SamPredictor(modelo_sam)
    
    imagen = cv2.imread(ruta_imagen)
    imagen_rgb = cv2.cvtColor(imagen, cv2.COLOR_BGR2RGB)
    predictor.set_image(imagen_rgb)
    
    # Supongamos que se tiene una caja delimitadora proporcionada por GroundingDINO
    caja_delimitadora = np.array([100, 100, 300, 300])  # Ejemplo
    cajas_tensor = torch.tensor(caja_delimitadora, dtype=torch.float).to(dispositivo)
    cajas_transformadas = predictor.transform.apply_boxes_torch(cajas_tensor, imagen.shape[:2]).to(dispositivo)
    
    mascaras, _, _ = predictor.predict_torch(
        point_coords=None,
        point_labels=None,
        boxes=cajas_transformadas,
        multimask_output=False,
    )
    
    # Redimensionar máscara a tamaño objetivo
    objetivo_tamano = (480, 640)
    mascaras_redim = torch.nn.functional.interpolate(mascaras.float(), size=objetivo_tamano, mode="nearest")
    
    # Convertir a imagen binaria
    mascara_binaria = np.zeros((objetivo_tamano[0], objetivo_tamano[1]), dtype=np.uint8)
    mascara_binaria[mascaras_redim.cpu().numpy()[0, 0] > 0] = 255
    
    cv2.imwrite("mascara_primera.png", mascara_binaria)

if __name__ == "__main__":
    cargar_y_predecir_mascara("test.png", "objeto de interés", "sam_vit_b.pth")

3.4 Preparación de los archivos mesh

El modelo 3D se puede reconstruir mediante herramientas como BundleSDF o aplicaciones de escaneo 3D con cámaras de profundidad. Los archivos resultantes (OBJ, MTL y textura PNG) se colocan en la carpeta mesh.

  1. Ejecución del demo

Una vez preparado el conjunto de datos, se ejecuta el demo con el siguiente script modificado:

import os
import argparse
import numpy as np
import trimesh
from estimador import *  # Suponiendo módulos renombrados
from lector_datos import *

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    directorio_codigo = os.path.dirname(os.path.realpath(__file__))
    parser.add_argument('--archivo_malla', type=str, default=f'{directorio_codigo}/datos_demo/malla.obj')
    parser.add_argument('--directorio_escena', type=str, default=f'{directorio_codigo}/datos_demo')
    parser.add_argument('--iteraciones_estimacion', type=int, default=5)
    parser.add_argument('--iteraciones_rastreo', type=int, default=2)
    parser.add_argument('--modo_debug', type=int, default=1)
    parser.add_argument('--directorio_debug', type=str, default=f'{directorio_codigo}/depuracion')
    args = parser.parse_args()

    malla = trimesh.load(args.archivo_malla)
    origen, extensiones = trimesh.bounds.oriented_bounds(malla)
    caja = np.stack([-extensiones/2, extensiones/2], axis=0).reshape(2,3)

    estimador = PoseEstimator(pts_modelo=malla.vertices, normales_modelo=malla.vertex_normals, malla=malla, debug_dir=args.directorio_debug, debug=args.modo_debug)
    lector = LectorEscena(directorio_video=args.directorio_escena, lado_corto=None, zfar=np.inf)

    for indice in range(len(lector.archivos_color)):
        color = lector.obtener_color(indice)
        profundidad = lector.obtener_profundidad(indice)
        if indice == 0:
            mascara = lector.obtener_mascara(0).astype(bool)
            pose = estimador.registrar(K=lector.K, rgb=color, depth=profundidad, mascara_ob=mascara, iteraciones=args.iteraciones_estimacion)
        else:
            pose = estimador.rastrear_uno(rgb=color, depth=profundidad, K=lector.K, iteraciones=args.iteraciones_rastreo)

        os.makedirs(os.path.join(args.directorio_debug, 'poses'), exist_ok=True)
        np.savetxt(os.path.join(args.directorio_debug, 'poses', f'{lector.ids[indice]}.txt'), pose.reshape(4,4))

Ejecutar con el comando: python ejecutar_demo.py --modo_debug 2

  1. Próximos pasos

Se planea crear cnojuntos de entrenamiento para LineMod y desplegar el proyecto en dispositivos Jetson para apliccaiones de agarre robótico.

Etiquetas: FoundationPose IntelRealSense opencv Python 3DReconstruction

Publicado el 5-31 21:57