Enhancing MedSAM with a Lightweight Box Predictor for Medical Image Segmentation

Resumen Técnico: Mejora de MedSAM con un Predictor de Cajas Ligero para la Segmentación de Imágenes Médicas

Planteamiento del Problema
La segmentación semántica en el diagnóstico por imágenes médicas se ve obstaculizada por la escasez de datos, la alta variabilidad entre modalidades (TC, RM, ecografía) y la dificultad de obtener conjuntos de datos anotados a gran escala. Aunque los modelos fundacionales como el Segment Anything Model (SAM) ofrecen capacidades de segmentación de propósito general, su aplicación directa a imágenes médicas es limitada debido a las brechas de dominio. Específicamente, la adaptación de MedSAM, aunque mejorada, sigue siendo altamente sensible a la calidad de los prompts del usuario. El prompt de punto estándar, aunque natural para la interacción del usuario, a menudo proporciona un contexto espacial insuficiente para una segmentación fiable, particularmente cuando las estructuras objetivo son irregulares, pequeñas o poseen bajo contraste. Esta ambigüedad conduce a una degradación significativa del rendimiento, especialmente cuando el modelo no está totalmente ajustado (fine-tuned) o cuando los clics del usuario son imprecisos.

Metodología
Los autores proponen un marco de segmentación mejorado que integra un módulo de Predictor de Cajas ligero en la arquitectura de MedSAM. La innovación central es la conversión de un único clic de usuario (prompt de punto) en una caja delimitadora aproximada, que sirve como un prior espacial robusto para guiar el proceso de segmentación.

El método opera mediante un pipeline de entrenamiento de dos etapas:

1. Etapa 1 (Entrenamiento del Predictor de Cajas): Se entrena un Perceptrón Multicapa (MLP) ligero de forma independiente para predecir los parámetros de la caja delimitadora ($W, H, \Delta x, \Delta y$) a partir de un parche de incrustaciones (embeddings) de imagen localizado de $5 \times 5$ extraído alrededor del punto de entrada. El punto de entrada está sujeto a una perturbación espacial aleatoria durante el entrenamiento (escalada al 30% del área de la máscara promedio) para simular la inexactitud realista del usuario. El módulo se optimiza utilizando una función de pérdida compuesta que combina la pérdida $L1$ y la pérdida de Intersección sobre Unión Generalizada (GIoU).
2. Etapa 2 (Integración y Ajuste Fino): El Predictor de Cajas preentrenado se integra en el pipeline de MedSAM. El Codificador de Imágenes y el Predictor de Cajas se mantienen congelados para preservar las representaciones aprendidas. El Codificador de Prompts y el Decodificador de Máscaras se ajustan opcionalmente. La caja delimitadora predicha se introduce en el Codificador de Prompts junto con el prompt de punto original (o de forma independiente) para generar la máscara de segmentación final.

El Predictor de Cajas añade solo 1.6 millones de parámetros e introduce un overhead de inferencia insignificante (aproximadamente 61–63 ms por imagen en una RTX 4090).

Contribuciones Clave

• Predictor de Cajas Ligero: Un módulo que transforma un prompt de punto único en una caja delimitadora informativa, proporcionando una guía espacial que estabiliza MedSAM, particularmente cuando los prompts de punto por sí solos son ambiguos o están mal ubicados.
• Pipeline de Entrenamiento de Dos Etapas: Una estrategia donde el Predictor de Cajas se entrena por separado y luego se integra como una guía espacial fija. Esto permite que el sistema logre una segmentación robusta incluso cuando el modelo MedSAM principal permanece congelado, lo que lo hace adecuado para entornos con recursos limitados.
• Evaluación Sistemática: Evaluación exhaustiva a través de cuatro conjuntos de datos heterogéneos (FLARE22, BRISC, BUSI, LungSegDB) que cubren las modalidades de TC, RM y ecografía. El estudio evalúa seis configuraciones, analizando la interacción entre los tipos de prompt (punto, caja, híbrido) y los estados del modelo (congelado vs. ajustado).
• Análisis de Robustez de Prompts: Investigación sobre cómo las perturbaciones de entrada y los ajustes del tamaño de la caja afectan la precisión, demostrando que el método sigue siendo efectivo incluso con clics descentrados y requiere un ajuste manual mínimo para diferentes modalidades.

Resultados Experimentales
El método propuesto se evaluó frente a los baselines de MedSAM y otras variantes de SAM (SAM2, SAM2UNet) a través de cuatro conjuntos de datos. Los hallazgos clave incluyen:

• Ganancias de Rendimiento: El método mejoró consistentemente la precisión de la segmentación. En el conjunto de datos BUSI, la puntuación Dice aumentó de 0.8771 (MedSAM + punto) a 0.8904 (MedSAM + Caja). En FLARE22, la puntuación Dice mejoró de 0.9169 a 0.9314.
• Recuperación de Modelo Congelado: En un entorno "congelado" donde los pesos de MedSAM no fueron actualizados, depender únicamente de los prompts de punto resultó en un fallo casi total (puntuaciones Dice tan bajas como 1.64% en LungSegDB). Sin embargo, la introducción de la caja delimitadora predicha restauró significativamente el rendimiento (por ejemplo, el Dice subió a 72.58% en LungSegDB), demostrando que el prior geométrico es crítico para la operación del modelo.
• Adaptación de Modalidad: El método mostró comportamientos distintos según la modalidad. Los conjuntos de datos de ecografía y RM (BUSI, BRISC) se beneficiaron de un ligero agrandamiento de la caja delimitadora (10–20%) para compensar los límites ambiguos, mientras que los conjuntos de datos de TC (FLARE22, LungSegDB) funcionaron de manera óptima con las cajas predichas iniciales debido a su mayor contraste.
• Eficiencia: El enfoque logró estas ganancias con un costo computacional mínimo, manteniendo tiempos de inferencia comparables al baseline.

Significancia y Reivindicaciones
El artículo afirma que este trabajo no redefine el paradigma fundamental de la segmentación, sino que ofrece una mejora práctica para MedSAM. Su principal significancia radica en reducir la dependencia de prompts manuales precisos, mejorando así la robustez de la segmentación interactiva en flujos de trabajo clínicos reales. Al aprovechar un módulo auxiliar ligero, el método cierra la brecha entre las expectativas geométricas del modelo fundacional y las limitaciones del input disperso del usuario. Los autores enfatizan que este enfoque es particularmente valioso en entornos con datos etiquetados o recursos computacionales limitados, ya que permite una segmentación de alto rendimiento sin la necesidad de un ajuste fino completo del modelo. El estudio concluye que, si bien el método tiene limitaciones respecto a regiones múltiples y desconectadas, aborda eficazmente la ambigüedad de los prompts de punto en diversos dominios de imágenes médicas.