DCTFFTW::DCTFFTW(int _sizex, int _sizey, int _dctmode, int _bitsPerSample) :
DCTClass(_sizex , _sizey, _dctmode, _bitsPerSample)
{
int size2d = sizey*sizex;
int cursize = 1;
dctshift = 0;
while (cursize < size2d)
{
dctshift++;
cursize = (cursize<<1);
}
dctshift0 = dctshift + 2;
fSrc = (float *)fftwf_malloc(sizeof(float) * size2d );
fSrcDCT = (float *)fftwf_malloc(sizeof(float) * size2d );
g_fftw_plans_mutex.lock();
dctplan = fftwf_plan_r2r_2d(sizey, sizex, fSrc, fSrcDCT,
FFTW_REDFT10, FFTW_REDFT10, FFTW_ESTIMATE); // direct fft
g_fftw_plans_mutex.unlock();
}
Image FFTPoisson::apply(Window dx, Window dy, Window target, float targetStrength) {
assert(dx.width == dy.width &&
dx.height == dy.height &&
dx.frames == dy.frames &&
dx.channels == dy.channels,
"x gradient must be same size as y gradient\n");
if (target) {
assert(target.width == dx.width &&
target.height == dx.height &&
target.frames == dx.frames &&
target.channels == dx.channels,
"target image must have the same size as the gradient images\n");
}
Image fftBuff(dx.width, dx.height, dx.frames, 1);
//compute two 1D lookup tables for computing the DCT of a 2D Laplacian on the fly
Image ftLapY(1, dx.height, 1, 1);
Image ftLapX(dx.width, 1, 1, 1);
for(int x = 0; x < dx.width; x++)
ftLapX(x, 0)[0] = 2.0f * cos((M_PI * x) / (dx.width - 1));
for(int y = 0; y < dx.height; y++)
ftLapY(0, y)[0] = -4.0f + (2.0f * cos((M_PI * y) / (dx.height - 1)));
// Create a DCT-I plan, which is its own inverse.
fftwf_plan fftPlan;
fftPlan = fftwf_plan_r2r_2d(dx.height, dx.width,
fftBuff(0, 0), fftBuff(0, 0),
FFTW_REDFT00, FFTW_REDFT00, FFTW_ESTIMATE); //use FFTW_PATIENT when plan can be reused
Image out(dx.width, dx.height, dx.frames, dx.channels);
for (int t = 0; t < dx.frames; t++) {
for (int c = 0; c < dx.channels; c++) {
float dcSum = 0.0f;
// compute h_hat from u, gx, gy (see equation 48 in the paper), as well as the DC term of u's DCT.
float *fftPtr = fftBuff(0, 0);
for(int y = 0; y < dx.height; y++) {
for(int x = 0; x < dx.width; x++) {
// Compute DC term of u's DCT without computing the whole DCT.
float dcMult = 1.0f;
if ((x > 0) && (x < dx.width - 1))
dcMult *= 2.0f;
if ((y > 0) && (y < dx.height - 1))
dcMult *= 2.0f;
if (target) {
dcSum += dcMult * target(x, y, t)[c];
} else {
// try to read the dc term out of the double
// integral of the gradient fields
// instead. Works if the gradients were
// computed with a zero boundary condition.
dcSum += 2.0f*((dx.width-x)*dx(x, y, t)[c] + (dy.height-y)*dy(x, y, t)[c]);
}
if (target)
*fftPtr = targetStrength * target(x, y, t)[c];
else
*fftPtr = 0;
// Subtract g^x_x and g^y_y, with boundary factor of -2.0 to account for boundary reflections implicit in the DCT
if (x == 0) {
*fftPtr -= (+2.0f * dx(x+1, y, t)[c]);
} else if (x == dx.width - 1) {
*fftPtr -= (-2.0f * dx(x, y, t)[c]);
} else {
*fftPtr -= (dx(x+1, y, t)[c] - dx(x, y, t)[c]);
}
if (y == 0) {
*fftPtr -= (+2.0f * dy(x, y+1, t)[c]);
} else if (y == dx.height -1) {
*fftPtr -= (-2.0f * dy(x, y, t)[c]);
} else {
*fftPtr -= (dy(x, y+1, t)[c] - dy(x, y, t)[c]);
}
fftPtr++;
}
}
//transform h_hat to H_hat by taking the DCT of h_hat
fftwf_execute(fftPlan);
//compute F_hat using H_hat (see equation 29 in the paper)
fftPtr = fftBuff(0, 0);
for(int y = 0; y < dx.height; y++) {
for(int x = 0; x < dx.width; x++) {
float ftLapResponse = ftLapY(0, y)[0] + ftLapX(x, 0)[0];
*fftPtr++ /= (targetStrength - ftLapResponse);
}
}
fftBuff(0, 0)[0] = dcSum;
//transform F_hat to f_hat by taking the inverse DCT of F_hat
fftwf_execute(fftPlan);
float fftMult = 1.0f / (4.0f * (dx.width-1) * (dx.height-1));
fftPtr = fftBuff(0, 0);
for(int y = 0; y < dx.height; y++) {
for(int x = 0; x < dx.width; x++) {
out(x, y, t)[c] = (*fftPtr++) * fftMult;
}
}
}
}
fftwf_destroy_plan(fftPlan);
return out;
}
