elbg.c 13.5 KB
Newer Older
1
/*
2
 * Copyright (C) 2007 Vitor Sessak <vitor1001@gmail.com>
3
 *
4
 * This file is part of Libav.
5
 *
6
 * Libav is free software; you can redistribute it and/or
7 8 9 10
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
11
 * Libav is distributed in the hope that it will be useful,
12 13 14 15 16
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
17
 * License along with Libav; if not, write to the Free Software
18 19 20 21
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

/**
22
 * @file
23 24 25 26 27
 * Codebook Generator using the ELBG algorithm
 */

#include <string.h>

28
#include "libavutil/common.h"
29
#include "libavutil/lfg.h"
30 31 32
#include "elbg.h"
#include "avcodec.h"

33
#define DELTA_ERR_MAX 0.1  ///< Precision of the ELBG algorithm (as percentage error)
34 35 36 37 38 39 40 41 42 43 44 45

/**
 * In the ELBG jargon, a cell is the set of points that are closest to a
 * codebook entry. Not to be confused with a RoQ Video cell. */
typedef struct cell_s {
    int index;
    struct cell_s *next;
} cell;

/**
 * ELBG internal data
 */
46
typedef struct elbg_data {
47 48 49 50 51 52 53 54 55
    int error;
    int dim;
    int numCB;
    int *codebook;
    cell **cells;
    int *utility;
    int *utility_inc;
    int *nearest_cb;
    int *points;
56
    AVLFG *rand_state;
Måns Rullgård's avatar
Måns Rullgård committed
57
    int *scratchbuf;
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
} elbg_data;

static inline int distance_limited(int *a, int *b, int dim, int limit)
{
    int i, dist=0;
    for (i=0; i<dim; i++) {
        dist += (a[i] - b[i])*(a[i] - b[i]);
        if (dist > limit)
            return INT_MAX;
    }

    return dist;
}

static inline void vect_division(int *res, int *vect, int div, int dim)
{
    int i;
    if (div > 1)
        for (i=0; i<dim; i++)
            res[i] = ROUNDED_DIV(vect[i],div);
    else if (res != vect)
        memcpy(res, vect, dim*sizeof(int));

}

static int eval_error_cell(elbg_data *elbg, int *centroid, cell *cells)
{
    int error=0;
    for (; cells; cells=cells->next)
        error += distance_limited(centroid, elbg->points + cells->index*elbg->dim, elbg->dim, INT_MAX);

    return error;
}

static int get_closest_codebook(elbg_data *elbg, int index)
{
    int i, pick=0, diff, diff_min = INT_MAX;
    for (i=0; i<elbg->numCB; i++)
        if (i != index) {
            diff = distance_limited(elbg->codebook + i*elbg->dim, elbg->codebook + index*elbg->dim, elbg->dim, diff_min);
            if (diff < diff_min) {
                pick = i;
                diff_min = diff;
            }
        }
    return pick;
}

static int get_high_utility_cell(elbg_data *elbg)
{
    int i=0;
    /* Using linear search, do binary if it ever turns to be speed critical */
110
    int r = av_lfg_get(elbg->rand_state)%elbg->utility_inc[elbg->numCB-1] + 1;
111 112
    while (elbg->utility_inc[i] < r)
        i++;
113

114
    assert(elbg->cells[i]);
115

116 117 118 119 120 121
    return i;
}

/**
 * Implementation of the simple LBG algorithm for just two codebooks
 */
Måns Rullgård's avatar
Måns Rullgård committed
122 123
static int simple_lbg(elbg_data *elbg,
                      int dim,
124
                      int *centroid[3],
125 126 127 128 129 130
                      int newutility[3],
                      int *points,
                      cell *cells)
{
    int i, idx;
    int numpoints[2] = {0,0};
Måns Rullgård's avatar
Måns Rullgård committed
131 132 133 134
    int *newcentroid[2] = {
        elbg->scratchbuf + 3*dim,
        elbg->scratchbuf + 4*dim
    };
135 136
    cell *tempcell;

Måns Rullgård's avatar
Måns Rullgård committed
137
    memset(newcentroid[0], 0, 2 * dim * sizeof(*newcentroid[0]));
138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166

    newutility[0] =
    newutility[1] = 0;

    for (tempcell = cells; tempcell; tempcell=tempcell->next) {
        idx = distance_limited(centroid[0], points + tempcell->index*dim, dim, INT_MAX)>=
              distance_limited(centroid[1], points + tempcell->index*dim, dim, INT_MAX);
        numpoints[idx]++;
        for (i=0; i<dim; i++)
            newcentroid[idx][i] += points[tempcell->index*dim + i];
    }

    vect_division(centroid[0], newcentroid[0], numpoints[0], dim);
    vect_division(centroid[1], newcentroid[1], numpoints[1], dim);

    for (tempcell = cells; tempcell; tempcell=tempcell->next) {
        int dist[2] = {distance_limited(centroid[0], points + tempcell->index*dim, dim, INT_MAX),
                       distance_limited(centroid[1], points + tempcell->index*dim, dim, INT_MAX)};
        int idx = dist[0] > dist[1];
        newutility[idx] += dist[idx];
    }

    return newutility[0] + newutility[1];
}

static void get_new_centroids(elbg_data *elbg, int huc, int *newcentroid_i,
                              int *newcentroid_p)
{
    cell *tempcell;
Måns Rullgård's avatar
Måns Rullgård committed
167 168
    int *min = newcentroid_i;
    int *max = newcentroid_p;
169 170 171 172 173 174 175 176 177 178 179 180 181 182
    int i;

    for (i=0; i< elbg->dim; i++) {
        min[i]=INT_MAX;
        max[i]=0;
    }

    for (tempcell = elbg->cells[huc]; tempcell; tempcell = tempcell->next)
        for(i=0; i<elbg->dim; i++) {
            min[i]=FFMIN(min[i], elbg->points[tempcell->index*elbg->dim + i]);
            max[i]=FFMAX(max[i], elbg->points[tempcell->index*elbg->dim + i]);
        }

    for (i=0; i<elbg->dim; i++) {
Måns Rullgård's avatar
Måns Rullgård committed
183 184 185 186
        int ni = min[i] + (max[i] - min[i])/3;
        int np = min[i] + (2*(max[i] - min[i]))/3;
        newcentroid_i[i] = ni;
        newcentroid_p[i] = np;
187 188 189 190 191
    }
}

/**
 * Add the points in the low utility cell to its closest cell. Split the high
192
 * utility cell, putting the separated points in the (now empty) low utility
193 194 195 196 197 198 199
 * cell.
 *
 * @param elbg         Internal elbg data
 * @param indexes      {luc, huc, cluc}
 * @param newcentroid  A vector with the position of the new centroids
 */
static void shift_codebook(elbg_data *elbg, int *indexes,
200
                           int *newcentroid[3])
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252
{
    cell *tempdata;
    cell **pp = &elbg->cells[indexes[2]];

    while(*pp)
        pp= &(*pp)->next;

    *pp = elbg->cells[indexes[0]];

    elbg->cells[indexes[0]] = NULL;
    tempdata = elbg->cells[indexes[1]];
    elbg->cells[indexes[1]] = NULL;

    while(tempdata) {
        cell *tempcell2 = tempdata->next;
        int idx = distance_limited(elbg->points + tempdata->index*elbg->dim,
                           newcentroid[0], elbg->dim, INT_MAX) >
                  distance_limited(elbg->points + tempdata->index*elbg->dim,
                           newcentroid[1], elbg->dim, INT_MAX);

        tempdata->next = elbg->cells[indexes[idx]];
        elbg->cells[indexes[idx]] = tempdata;
        tempdata = tempcell2;
    }
}

static void evaluate_utility_inc(elbg_data *elbg)
{
    int i, inc=0;

    for (i=0; i < elbg->numCB; i++) {
        if (elbg->numCB*elbg->utility[i] > elbg->error)
            inc += elbg->utility[i];
        elbg->utility_inc[i] = inc;
    }
}


static void update_utility_and_n_cb(elbg_data *elbg, int idx, int newutility)
{
    cell *tempcell;

    elbg->utility[idx] = newutility;
    for (tempcell=elbg->cells[idx]; tempcell; tempcell=tempcell->next)
        elbg->nearest_cb[tempcell->index] = idx;
}

/**
 * Evaluate if a shift lower the error. If it does, call shift_codebooks
 * and update elbg->error, elbg->utility and elbg->nearest_cb.
 *
 * @param elbg  Internal elbg data
253
 * @param idx   {luc (low utility cell, huc (high utility cell), cluc (closest cell to low utility cell)}
254 255 256 257 258
 */
static void try_shift_candidate(elbg_data *elbg, int idx[3])
{
    int j, k, olderror=0, newerror, cont=0;
    int newutility[3];
Måns Rullgård's avatar
Måns Rullgård committed
259 260 261 262 263
    int *newcentroid[3] = {
        elbg->scratchbuf,
        elbg->scratchbuf + elbg->dim,
        elbg->scratchbuf + 2*elbg->dim
    };
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
    cell *tempcell;

    for (j=0; j<3; j++)
        olderror += elbg->utility[idx[j]];

    memset(newcentroid[2], 0, elbg->dim*sizeof(int));

    for (k=0; k<2; k++)
        for (tempcell=elbg->cells[idx[2*k]]; tempcell; tempcell=tempcell->next) {
            cont++;
            for (j=0; j<elbg->dim; j++)
                newcentroid[2][j] += elbg->points[tempcell->index*elbg->dim + j];
        }

    vect_division(newcentroid[2], newcentroid[2], cont, elbg->dim);

    get_new_centroids(elbg, idx[1], newcentroid[0], newcentroid[1]);

    newutility[2]  = eval_error_cell(elbg, newcentroid[2], elbg->cells[idx[0]]);
    newutility[2] += eval_error_cell(elbg, newcentroid[2], elbg->cells[idx[2]]);

    newerror = newutility[2];

Måns Rullgård's avatar
Måns Rullgård committed
287
    newerror += simple_lbg(elbg, elbg->dim, newcentroid, newutility, elbg->points,
288 289 290
                           elbg->cells[idx[1]]);

    if (olderror > newerror) {
Måns Rullgård's avatar
Måns Rullgård committed
291
        shift_codebook(elbg, idx, newcentroid);
292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315

        elbg->error += newerror - olderror;

        for (j=0; j<3; j++)
            update_utility_and_n_cb(elbg, idx[j], newutility[j]);

        evaluate_utility_inc(elbg);
    }
 }

/**
 * Implementation of the ELBG block
 */
static void do_shiftings(elbg_data *elbg)
{
    int idx[3];

    evaluate_utility_inc(elbg);

    for (idx[0]=0; idx[0] < elbg->numCB; idx[0]++)
        if (elbg->numCB*elbg->utility[idx[0]] < elbg->error) {
            if (elbg->utility_inc[elbg->numCB-1] == 0)
                return;

316
            idx[1] = get_high_utility_cell(elbg);
317 318
            idx[2] = get_closest_codebook(elbg, idx[0]);

319 320
            if (idx[1] != idx[0] && idx[1] != idx[2])
                try_shift_candidate(elbg, idx);
321 322 323 324 325
        }
}

#define BIG_PRIME 433494437LL

326 327 328
int ff_init_elbg(int *points, int dim, int numpoints, int *codebook,
                 int numCB, int max_steps, int *closest_cb,
                 AVLFG *rand_state)
329
{
330
    int i, k, ret = 0;
331 332 333 334 335

    if (numpoints > 24*numCB) {
        /* ELBG is very costly for a big number of points. So if we have a lot
           of them, get a good initial codebook to save on iterations       */
        int *temp_points = av_malloc(dim*(numpoints/8)*sizeof(int));
336 337
        if (!temp_points)
            return AVERROR(ENOMEM);
338 339 340 341 342
        for (i=0; i<numpoints/8; i++) {
            k = (i*BIG_PRIME) % numpoints;
            memcpy(temp_points + i*dim, points + k*dim, dim*sizeof(int));
        }

343 344 345 346 347 348 349 350
        ret = ff_init_elbg(temp_points, dim, numpoints / 8, codebook,
                           numCB, 2 * max_steps, closest_cb, rand_state);
        if (ret < 0) {
            av_freep(&temp_points);
            return ret;
        }
        ret = ff_do_elbg(temp_points, dim, numpoints / 8, codebook,
                         numCB, 2 * max_steps, closest_cb, rand_state);
351 352 353 354 355 356
        av_free(temp_points);

    } else  // If not, initialize the codebook with random positions
        for (i=0; i < numCB; i++)
            memcpy(codebook + i*dim, points + ((i*BIG_PRIME)%numpoints)*dim,
                   dim*sizeof(int));
357
    return ret;
358 359
}

360
int ff_do_elbg(int *points, int dim, int numpoints, int *codebook,
361
                int numCB, int max_steps, int *closest_cb,
362
                AVLFG *rand_state)
363 364 365 366
{
    int dist;
    elbg_data elbg_d;
    elbg_data *elbg = &elbg_d;
367
    int i, j, k, last_error, steps = 0, ret = 0;
368 369 370 371
    int *dist_cb = av_malloc(numpoints*sizeof(int));
    int *size_part = av_malloc(numCB*sizeof(int));
    cell *list_buffer = av_malloc(numpoints*sizeof(cell));
    cell *free_cells;
372
    int best_dist, best_idx = 0;
373 374 375 376 377 378 379 380 381 382

    elbg->error = INT_MAX;
    elbg->dim = dim;
    elbg->numCB = numCB;
    elbg->codebook = codebook;
    elbg->cells = av_malloc(numCB*sizeof(cell *));
    elbg->utility = av_malloc(numCB*sizeof(int));
    elbg->nearest_cb = closest_cb;
    elbg->points = points;
    elbg->utility_inc = av_malloc(numCB*sizeof(int));
Måns Rullgård's avatar
Måns Rullgård committed
383
    elbg->scratchbuf = av_malloc(5*dim*sizeof(int));
384

385 386 387 388 389 390
    if (!dist_cb || !size_part || !list_buffer || !elbg->cells ||
        !elbg->utility || !elbg->utility_inc || !elbg->scratchbuf) {
        ret = AVERROR(ENOMEM);
        goto out;
    }

391 392 393 394 395 396 397 398 399 400 401 402 403 404
    elbg->rand_state = rand_state;

    do {
        free_cells = list_buffer;
        last_error = elbg->error;
        steps++;
        memset(elbg->utility, 0, numCB*sizeof(int));
        memset(elbg->cells, 0, numCB*sizeof(cell *));

        elbg->error = 0;

        /* This loop evaluate the actual Voronoi partition. It is the most
           costly part of the algorithm. */
        for (i=0; i < numpoints; i++) {
405
            best_dist = distance_limited(elbg->points + i*elbg->dim, elbg->codebook + best_idx*elbg->dim, dim, INT_MAX);
406
            for (k=0; k < elbg->numCB; k++) {
407 408 409 410
                dist = distance_limited(elbg->points + i*elbg->dim, elbg->codebook + k*elbg->dim, dim, best_dist);
                if (dist < best_dist) {
                    best_dist = dist;
                    best_idx = k;
411 412
                }
            }
413 414
            elbg->nearest_cb[i] = best_idx;
            dist_cb[i] = best_dist;
415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442
            elbg->error += dist_cb[i];
            elbg->utility[elbg->nearest_cb[i]] += dist_cb[i];
            free_cells->index = i;
            free_cells->next = elbg->cells[elbg->nearest_cb[i]];
            elbg->cells[elbg->nearest_cb[i]] = free_cells;
            free_cells++;
        }

        do_shiftings(elbg);

        memset(size_part, 0, numCB*sizeof(int));

        memset(elbg->codebook, 0, elbg->numCB*dim*sizeof(int));

        for (i=0; i < numpoints; i++) {
            size_part[elbg->nearest_cb[i]]++;
            for (j=0; j < elbg->dim; j++)
                elbg->codebook[elbg->nearest_cb[i]*elbg->dim + j] +=
                    elbg->points[i*elbg->dim + j];
        }

        for (i=0; i < elbg->numCB; i++)
            vect_division(elbg->codebook + i*elbg->dim,
                          elbg->codebook + i*elbg->dim, size_part[i], elbg->dim);

    } while(((last_error - elbg->error) > DELTA_ERR_MAX*elbg->error) &&
            (steps < max_steps));

443
out:
444 445 446 447 448 449
    av_free(dist_cb);
    av_free(size_part);
    av_free(elbg->utility);
    av_free(list_buffer);
    av_free(elbg->cells);
    av_free(elbg->utility_inc);
Måns Rullgård's avatar
Måns Rullgård committed
450
    av_free(elbg->scratchbuf);
451
    return ret;
452
}