dupeguru/core_pe/modules/block/block.c

/* Created By: Virgil Dupras
 * Created On: 2010-01-30
 * Copyright 2010 Hardcoded Software (http://www.hardcoded.net)
 */

#define PY_SSIZE_T_CLEAN
#include "Python.h"

/* avgdiff/maxdiff has been called with empty lists */
static PyObject *NoBlocksError;
/* avgdiff/maxdiff has been called with 2 block lists of different size. */
static PyObject *DifferentBlockCountError;

static inline int
max(int a, int b)
{
    return b > a ? b : a;
}

static inline int
min(int a, int b)
{
    return b < a ? b : a;
}

/* Create a tuple out of an array of integers. */
static PyObject*
inttuple(int numbers[], int count)
{
    int i;
    PyObject *pnumber;
    PyObject *result;
    
    result = PyTuple_New(count);
    
    for (i=0; i<count; i++) {
        pnumber = PyInt_FromLong(numbers[i]);
        if (pnumber == NULL) {
            Py_DECREF(result);
            return NULL;
        }
        PyTuple_SET_ITEM(result, i, pnumber);
    }
    
    return result;
}

/* Returns a 3 sized tuple containing the mean color of 'image'.    
 * image: a PIL image or crop.
 */
static PyObject*
getblock(PyObject *image)
{
    int i, totals[3] = {0, 0, 0};
    Py_ssize_t pixel_count;
    PyObject *ppixels;
    
    ppixels = PyObject_CallMethod(image, "getdata", NULL);
    if (ppixels == NULL) {
        return NULL;
    }
    
    pixel_count = PySequence_Length(ppixels);
    for (i=0; i<pixel_count; i++) {
        PyObject *ppixel, *pr, *pg, *pb;
        int r, g, b;
        
        ppixel = PySequence_ITEM(ppixels, i);
        pr = PySequence_ITEM(ppixel, 0);
        pg = PySequence_ITEM(ppixel, 1);
        pb = PySequence_ITEM(ppixel, 2);
        Py_DECREF(ppixel);
        r = PyInt_AsSsize_t(pr);
        g = PyInt_AsSsize_t(pg);
        b = PyInt_AsSsize_t(pb);
        Py_DECREF(pr);
        Py_DECREF(pg);
        Py_DECREF(pb);
        
        totals[0] += r;
        totals[1] += g;
        totals[2] += b;
    }
    
    Py_DECREF(ppixels);
    
    if (pixel_count) {
        totals[0] /= pixel_count;
        totals[1] /= pixel_count;
        totals[2] /= pixel_count;
    }
    
    return inttuple(totals, 3);
}

/* Returns the difference between the first block and the second.
 * It returns an absolute sum of the 3 differences (RGB).
 */
static inline int
diff(PyObject *first, PyObject *second)
{
    Py_ssize_t r1, g1, b1, r2, b2, g2;
    PyObject *pr, *pg, *pb;
    pr = PySequence_ITEM(first, 0);
    pg = PySequence_ITEM(first, 1);
    pb = PySequence_ITEM(first, 2);
    r1 = PyInt_AsSsize_t(pr);
    g1 = PyInt_AsSsize_t(pg);
    b1 = PyInt_AsSsize_t(pb);
    Py_DECREF(pr);
    Py_DECREF(pg);
    Py_DECREF(pb);
    
    pr = PySequence_ITEM(second, 0);
    pg = PySequence_ITEM(second, 1);
    pb = PySequence_ITEM(second, 2);
    r2 = PyInt_AsSsize_t(pr);
    g2 = PyInt_AsSsize_t(pg);
    b2 = PyInt_AsSsize_t(pb);
    Py_DECREF(pr);
    Py_DECREF(pg);
    Py_DECREF(pb);
    
    return abs(r1 - r2) + abs(g1 - g2) + abs(b1 - b2);
}

PyDoc_STRVAR(block_getblocks2_doc,
"Returns a list of blocks (3 sized tuples).\n\
\n\
image: A PIL image to base the blocks on.\n\
block_count_per_side: This integer determine the number of blocks the function will return.\n\
If it is 10, for example, 100 blocks will be returns (10 width, 10 height). The blocks will not\n\
necessarely cover square areas. The area covered by each block will be proportional to the image\n\
itself.\n");

static PyObject*
block_getblocks2(PyObject *self, PyObject *args)
{
    int block_count_per_side, width, height, block_width, block_height, ih;
    PyObject *image;
    PyObject *pimage_size, *pwidth, *pheight;
    PyObject *result;
    
    if (!PyArg_ParseTuple(args, "Oi", &image, &block_count_per_side)) {
        return NULL;
    }
    
    pimage_size = PyObject_GetAttrString(image, "size");
    pwidth = PySequence_ITEM(pimage_size, 0);
    pheight = PySequence_ITEM(pimage_size, 1);
    width = PyInt_AsSsize_t(pwidth);
    height = PyInt_AsSsize_t(pheight);
    Py_DECREF(pimage_size);
    Py_DECREF(pwidth);
    Py_DECREF(pheight);
    
    if (!(width && height)) {
        return PyList_New(0);
    }
    
    block_width = max(width / block_count_per_side, 1);
    block_height = max(height / block_count_per_side, 1);
    
    result = PyList_New(block_count_per_side * block_count_per_side);
    if (result == NULL) {
        return NULL;
    }
    
    for (ih=0; ih<block_count_per_side; ih++) {
        int top, bottom, iw;
        top = min(ih*block_height, height-block_height);
        bottom = top + block_height;
        for (iw=0; iw<block_count_per_side; iw++) {
            int left, right, box[4];
            PyObject *pbox;
            PyObject *pmethodname;
            PyObject *pcrop;
            PyObject *pblock;
            
            left = min(iw*block_width, width-block_width);
            right = left + block_width;
            box[0] = left;
            box[1] = top;
            box[2] = right;
            box[3] = bottom;
            pbox = inttuple(box, 4);
            pmethodname = PyString_FromString("crop");
            pcrop = PyObject_CallMethodObjArgs(image, pmethodname, pbox);
            Py_DECREF(pmethodname);
            Py_DECREF(pbox);
            if (pcrop == NULL) {
                Py_DECREF(result);
                return NULL;
            }
            pblock = getblock(pcrop);
            Py_DECREF(pcrop);
            if (pblock == NULL) {
                Py_DECREF(result);
                return NULL;
            }
            PyList_SET_ITEM(result, ih*block_count_per_side+iw, pblock);
        }
    }
    
    return result;
}

PyDoc_STRVAR(block_avgdiff_doc,
"Returns the average diff between first blocks and seconds.\n\
\n\
If the result surpasses limit, limit + 1 is returned, except if less than min_iterations\n\
iterations have been made in the blocks.\n");

static PyObject*
block_avgdiff(PyObject *self, PyObject *args)
{
    PyObject *first, *second;
    int limit, min_iterations;
    Py_ssize_t count;
    int sum, i, result;
    
    if (!PyArg_ParseTuple(args, "OOii", &first, &second, &limit, &min_iterations)) {
        return NULL;
    }
    
    count = PySequence_Length(first);
    if (count != PySequence_Length(second)) {
        PyErr_SetString(DifferentBlockCountError, "");
        return NULL;
    }
    if (!count) {
        PyErr_SetString(NoBlocksError, "");
        return NULL;
    }
    
    sum = 0;
    for (i=0; i<count; i++) {
        int iteration_count;
        PyObject *item1, *item2;
        
        iteration_count = i + 1;
        item1 = PySequence_ITEM(first, i);
        item2 = PySequence_ITEM(second, i);
        sum += diff(item1, item2);
        Py_DECREF(item1);
        Py_DECREF(item2);
        if ((sum > limit*iteration_count) && (iteration_count >= min_iterations)) {
            return PyInt_FromSsize_t(limit + 1);
        }
    }
    
    result = sum / count;
    if (!result && sum) {
        result = 1;
    }
    return PyInt_FromSsize_t(result);
}

static PyMethodDef BlockMethods[] = {
    {"getblocks2",  block_getblocks2, METH_VARARGS, block_getblocks2_doc},
    {"avgdiff",  block_avgdiff, METH_VARARGS, block_avgdiff_doc},
    {NULL, NULL, 0, NULL} /* Sentinel */
};

PyMODINIT_FUNC
init_block(void)
{
    PyObject *m = Py_InitModule("_block", BlockMethods);
    if (m == NULL) {
        return;
    }
    
    NoBlocksError = PyErr_NewException("_block.NoBlocksError", NULL, NULL);
    PyModule_AddObject(m, "NoBlocksError", NoBlocksError);
    DifferentBlockCountError = PyErr_NewException("_block.DifferentBlockCountError", NULL, NULL);
    PyModule_AddObject(m, "DifferentBlockCountError", DifferentBlockCountError);
}