Haciendo a Python más rápido con Numba

Python es un lenguaje muy querido por su facilidad de uso y por la rapidez con la que permite escribir código. Sin embargo, una de sus limitaciones más conocidas es el rendimiento en tareas de cómputo intensivo, pues al tratarse de un lenguaje interpretado, suele ser más lento que uno compilado.

Esta limitación se hizo muy evidente mientras desarrollaba mi propia librería de algoritmos de procesamiento de imágenes que inicie justo en Python para aprovechar el amplio soporte de bibliotecas que tiene. En un inicio, la librería surgió únicamente como apoyo para los ejemplos de código de mis artículos, pero con el tiempo fue creciendo y comencé a utilizarla también en otros proyectos.

Aunque el rendimiento no siempre era el mejor, durante bastante tiempo no le di importancia, hasta que los tiempos de ejecución comenzaron a volverse un problema real, pero la cantidad de código y la complejidad hacían poco viable una reescritura completa en C++.

En lugar de optar por la reescritura, decidí explorar alternativas de optimización dentro del propio ecosistema de Python, comenzando por la simplificación de los algoritmos y la utilización de librerías para ciertas operaciones, pero la mejora más significativa que obtuve fue el uso de Numba.

Qué es Numba

Numba es un compilador just-in-time (JIT) para Python que transforma funciones numéricas en código máquina optimizado durante la ejecución. Su principal ventaja es acelerar secciones críticas del código sin abandonar Python ni reescribirlo en otro lenguaje.

El mayor beneficio se obtiene en funciones numéricas bien definidas, con control de flujo simple y uso intensivo de bucles, especialmente cuando operan sobre arreglos de NumPy con tipos homogéneos. En cambio, código fuertemente dependiente de objetos de Python, estructuras dinámicas o librerías ya optimizadas en C suele beneficiarse poco, e incluso puede verse afectado por la sobrecarga del proceso JIT.

Como ejemplo voy a mostrar el proceso de optimizacion de mi propia implementacion del algoritmo CLAHE. Al trabajar con cálculos locales y recorridos explícitos sobre los pixels de la imagen, resulta apropiado para mostrar cómo Numba puede acelerar este tipo de código en procesamiento de imágenes.

Breve explicación de qué es CLAHE

CLAHE (Contrast Limited Adaptive Histogram Equalization) es una variante del ecualizado de histograma que mejora el contraste de forma local. En lugar de aplicar una transformación global, divide la imagen en bloques y calcula un histograma independiente por región, lo que permite realzar detalles locales incluso en presencia de variaciones de iluminación.

Para evitar la amplificación del ruido, introduce un límite de recorte (clip limit) sobre los histogramas locales y redistribuye el exceso de forma controlada. Posteriormente, los resultados de cada bloque se combinan mediante interpolación para evitar discontinuidades visibles entre regiones.

Por su estructura basada en cálculos locales repetitivos y recorridos explícitos sobre los pixels de la imagen, hacen de CLAHE un buen candidato para analizar el impacto de distintas estrategias de optimización en Python.

En una implementación escrita solo en python, el coste de los bucles domina el tiempo de ejecución, dando lugar a un rendimiento muy inferior al de implementaciones optimizadas como la de OpenCV(hasta 17,000 veces mas lento), lo que hace evidente la necesidad de refactorizacion.

Identificación del cuello de botella

El primer paso en cualquier proceso de optimización consiste en identificar con claridad el cuello de botella, es decir, la operación o conjunto de operaciones que están dominando el tiempo total de ejecución. Para ello se pueden utilizar herramientas de análisis como cProfile, lo que permite obtener una visión global del comportamiento de la implementación del algoritmo.

print("=== cProfile: Custom CLAHE ===")
cProfile.run("HistogramEqualizationClaheGrayscale(img)")

El análisis muestra que la mayor parte del tiempo de ejecución no se concentraba en el cálculo de histogramas ni en la generación de las tablas de búsqueda, sino en la etapa de interpolación final. En particular, la función encargada de aplicar la interpolación espacial entre bloques acumulaba más del 60% del tiempo total, con cientos de miles de llamadas a funciones auxiliares y operaciones elementales dentro de bucles en Python.

Este resultado dejaba claro que el principal problema de rendimiento no esta en el algoritmo en sí, sino en la ejecución repetitiva de bucles y operaciones a bajo nivel dentro del intérprete de Python.

La interpolación como principal problema

Al inspeccionar el perfil, la función que más tiempo consume es la encargada de la interpolación bilineal:

1.096 s  ApplyInterpolation
0.234 s  InterpolationIndices

Es decir, la mayor parte del tiempo se recorre píxel a píxel en la imagen, combinando los mapas de transformación generados.

Implementación original: clara pero costosa

La versión original priorizaba la legibilidad y la separación de responsabilidades. Conceptualmente es correcta, pero desde el punto de vista del rendimiento es muy cara:

for y in range(height):
    i0, i1, y_weight = InterpolationIndices(y, cell_h, n_rows)

    for x in range(width):
        j0, j1, x_weight = InterpolationIndices(x, cell_w, n_cols)
        ...

Aquí ocurre lo peor para Python:

• Bucles anidados en Python puro

• Llamadas repetidas a InterpolationIndices

• Uso de np.floor y np.round en contexto escalar

Cada una de estas operaciones es pequeña, pero ejecutadas cientos de miles de veces se vuelven dominantes.

Enfoque optimizado: menos abstracción, más control

La versión optimizada cambia completamente el enfoque. En lugar de llamar a funciones auxiliares, toda la lógica de interpolación se integra directamente dentro del bucle, y se compila con Numba:

@njit(cache=True, fastmath=True)
def apply_interpolation_numba(...):
    for y in range(h):
        fy = (y + 0.5) / cell_h - 0.5
        ...
        for x in range(w):
            fx = (x + 0.5) / cell_w - 0.5
            ...

Las decisiones clave aquí son:

• Eliminar llamadas a funciones dentro del bucle crítico

• Reemplazar operaciones de NumPy por aritmética simple

• Usar tipos estáticos y bucles compilados con @njit

Esto permite que Numba genere código máquina eficiente, muy cercano a lo que escribirías en C o C++.

Este ejemplo deja una lección clara:
cuando el cuello de botella está en bucles intensivos y operaciones simples repetidas muchas veces, el problema no es Python como lenguaje, sino el coste del intérprete. En esos casos, Numba encaja de forma natural como una herramienta de optimización, especialmente en contextos de investigación donde la claridad del código sigue siendo importante. ahora las ordenes de magnitud estan solo en 5 y no en 170000.

Evitar comprobaciones repetitivas en el camino crítico

Una optimización menos evidente, pero igual de importante, es, eliminar validaciones y comprobaciones dentro de funciones que se ejecutan miles de veces. En Python solemos escribir código defensivo, validando tipos, dimensiones y rangos en cada llamada. Esto es correcto desde el punto de vista de la robustez, pero puede ser muy costoso cuando esas funciones forman parte del camino crítico del algoritmo.

Por ejemplo, en la versión original del cálculo del histograma se realizan varias comprobaciones en cada llamada:

def cal_histogram(channel):
    if not isinstance(channel, np.ndarray):
        raise TypeError(...)

    if channel.ndim != 2:
        raise ValueError(...)

    if channel.dtype != np.uint8:
        raise ValueError(...)

Estas validaciones son razonables si la función se usa de forma aislada o como API pública. El problema es que, en CLAHE, esta función se ejecuta una vez por bloque, y el número de bloques puede crecer rápidamente. El coste de estas comprobaciones termina siendo mayor que el propio cálculo del histograma.

Validar una vez, no cien veces

La clave aquí es separar responsabilidades:

• Las validaciones deben hacerse una sola vez, al inicio del procesamiento.

• Las funciones internas deben asumir que los datos ya son correctos.

En el contexto de CLAHE, tiene mucho más sentido comprobar que la imagen es uint8, 2D y válidar antes de entrar al bucle principal, y no dentro de cada llamada al histograma o al recorte del histograma.

Este cambio reduce drásticamente el trabajo innecesario del intérprete y, siguiendo esa misma idea, las versiones especializadas para el núcleo del algoritmo —como las optimizadas con Numba— eliminan por completo las comprobaciones para centrarse exclusivamente en el cálculo.

@njit(cache=True, fastmath=True)
def cal_histogram_numba(block):
    hist = np.zeros(256, dtype=np.int32)
    h, w = block.shape

    for y in range(h):
        for x in range(w):
            hist[block[y, x]] += 1

    return hist

Aquí se asume explícitamente que:

• block es 2D

• los valores están en [0, 255]

• el tipo de dato es el esperado

No hay verificaciones, no hay conversiones, no hay llamadas a funciones externas. Solo bucles y operaciones simples, exactamente lo que Numba puede optimizar bien.

Lo mismo ocurre con el recorte del histograma:

@njit(cache=True, fastmath=True)
def clip_histogram_numba(hist, clip_limit):
    excess = 0
    for i in range(256):
        if hist[i] > clip_limit:
            excess += hist[i] - clip_limit
            hist[i] = clip_limit

    incr = excess // 256
    for i in range(256):
        hist[i] += incr

    return hist

Ejemplos

conclusion

La idea central es simple: el código rápido no es el más defensivo, sino el que ejecuta solo lo necesario, una vez y en el lugar correcto. En algoritmos intensivos, como los de procesamiento de imágenes, conviene validar al inicio, fallar rápido y mantener el núcleo del algoritmo lo más simple y predecible posible.

Numba refuerza esta mentalidad al separar claramente la lógica de alto nivel del código crítico de bajo nivel. Aunque Python tenga límites en cómputo intensivo, un algoritmo bien diseñado —con lógica simplificada y sin trabajo innecesario en el camino crítico— puede acercarse mucho al rendimiento de código compilado usando herramientas como NumPy y Numba.