numpyの配列を受け取るCモジュールを作る

今作っているツールに次のようなコードがあるのだがかなり遅い。

# RGB配列とAlpha配列を合成してRGBAにするコード
image=wx.Image(path)
w=image.GetWidth()
h=image.GetHeight()
# RGB
rgb=image.GetData()
# Alphaチャンネル
alpha=image.GetAlphaData()

# 後でglTexImage2Dに使う
rgba=numpy.array([255]*w*h*4, <del>numpy.uint8</del>'u8')
rgb_index=0
alpha_index=0
rgba_index=0

# 合成する
for y in xrange(h):
    for x in xrange(w):
        rgba[rgba_index]=rgb[rgb_index]
        rgba[rgba_index+1]=rgb[rgb_index+1]
        rgba[rgba_index+2]=rgb[rgb_index+2]
        rgba[rgba_index+3]=alpha[alpha_index]
        rgba_index+=4
        rgb_index+=3
        alpha_index+=1

これをCで置き換えて処理時間の短縮を図ることにした。
以下のインターフェースにする。

imageutil.merge(w, h, rgba, rgb, alpha)

imageutil.c

#include <Python.h>
#include "Numeric/arrayobject.h"

// 処理本体
static void
merge(int w, int h, unsigned char *dst, 
    const unsigned char *rgb, const unsigned char *alpha)
{
  int x, y;
  for(y=0; y<h; ++y){
    for(x=0; x<w; ++x, dst+=4, rgb+=3, alpha+=1){
      dst[0]=rgb[0];
      dst[1]=rgb[1];
      dst[2]=rgb[2];
      dst[3]=*alpha;
    }
  }
}

// pythonの引数をC関数に渡す
static PyObject*
imageutil_merge(PyObject* self, PyObject* args)
{
  int w, h;
  // wx.Image.GetDataBuffer(), wx.Image.GetAlphaBuffer()を受け取る
  char *rgb, *alpha;
  // numpy.array([], numpy.uint8)を受け取る
  PyArrayObject *array;
  if (!PyArg_ParseTuple(args, "iiOww", 
        &w, &h, 
        &array, 
        &rgb, &alpha)){
    if (array->nd != 1 || array->descr->type_num != PyArray_UBYTE) {
      PyErr_SetString(PyExc_ValueError,
          "array must be unsigned char*");
      return NULL;

    }
  }
  merge(w, h, array->data, rgb, alpha);

  Py_RETURN_NONE;
}

// モジュールメソッド
static PyMethodDef imageutilmethods[] = {
  {"merge", imageutil_merge, METH_VARARGS},
  {NULL}, // sentinel
};

// モジュール登録
void initimageutil() {
  Py_InitModule("imageutil", imageutilmethods);
}

setup.py

from distutils.core import setup, Extension
setup(name="imageutil", version="1.0",
        ext_modules=[Extension("imageutil", ["imageutil.c"])])

ビルド

$ python setup.h build
# build/lib.linux-i686-2.6/imageutil.so が作成される

# debug symbol付き
$ python setup.h build -g

256x256ピクセルを想定した実行時間をtimeit(10)で計測してみたところ

2.6698679924 # python
0.00899910926819 # cモジュール

という結果に。
劇的に早くなった。
pixel毎の画像操作はさすがに苦手なのね。