FFmpeg: libavfilter/vf_deshake.c Source File

00001 /*
00002  * Copyright (C) 2010 Georg Martius <georg.martius@web.de>
00003  * Copyright (C) 2010 Daniel G. Taylor <dan@programmer-art.org>
00004  *
00005  * This file is part of FFmpeg.
00006  *
00007  * FFmpeg is free software; you can redistribute it and/or
00008  * modify it under the terms of the GNU Lesser General Public
00009  * License as published by the Free Software Foundation; either
00010  * version 2.1 of the License, or (at your option) any later version.
00011  *
00012  * FFmpeg is distributed in the hope that it will be useful,
00013  * but WITHOUT ANY WARRANTY; without even the implied warranty of
00014  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
00015  * Lesser General Public License for more details.
00016  *
00017  * You should have received a copy of the GNU Lesser General Public
00018  * License along with FFmpeg; if not, write to the Free Software
00019  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
00020  */
00021 
00052 #include "avfilter.h"
00053 #include "formats.h"
00054 #include "internal.h"
00055 #include "video.h"
00056 #include "libavutil/common.h"
00057 #include "libavutil/mem.h"
00058 #include "libavutil/pixdesc.h"
00059 #include "libavcodec/dsputil.h"
00060 
00061 #include "transform.h"
00062 
00063 #define CHROMA_WIDTH(link)  -((-link->w) >> av_pix_fmt_desc_get(link->format)->log2_chroma_w)
00064 #define CHROMA_HEIGHT(link) -((-link->h) >> av_pix_fmt_desc_get(link->format)->log2_chroma_h)
00065 
00066 enum SearchMethod {
00067     EXHAUSTIVE,        
00068     SMART_EXHAUSTIVE,  
00069     SEARCH_COUNT
00070 };
00071 
00072 typedef struct {
00073     int x;             
00074     int y;             
00075 } IntMotionVector;
00076 
00077 typedef struct {
00078     double x;             
00079     double y;             
00080 } MotionVector;
00081 
00082 typedef struct {
00083     MotionVector vector;  
00084     double angle;         
00085     double zoom;          
00086 } Transform;
00087 
00088 typedef struct {
00089     AVClass av_class;
00090     AVFilterBufferRef *ref;    
00091     int rx;                    
00092     int ry;                    
00093     enum FillMethod edge;      
00094     int blocksize;             
00095     int contrast;              
00096     enum SearchMethod search;  
00097     AVCodecContext *avctx;
00098     DSPContext c;              
00099     Transform last;            
00100     int refcount;              
00101     FILE *fp;
00102     Transform avg;
00103     int cw;                    
00104     int ch;
00105     int cx;
00106     int cy;
00107 } DeshakeContext;
00108 
00109 static int cmp(const double *a, const double *b)
00110 {
00111     return *a < *b ? -1 : ( *a > *b ? 1 : 0 );
00112 }
00113 
00117 static double clean_mean(double *values, int count)
00118 {
00119     double mean = 0;
00120     int cut = count / 5;
00121     int x;
00122 
00123     qsort(values, count, sizeof(double), (void*)cmp);
00124 
00125     for (x = cut; x < count - cut; x++) {
00126         mean += values[x];
00127     }
00128 
00129     return mean / (count - cut * 2);
00130 }
00131 
00138 static void find_block_motion(DeshakeContext *deshake, uint8_t *src1,
00139                               uint8_t *src2, int cx, int cy, int stride,
00140                               IntMotionVector *mv)
00141 {
00142     int x, y;
00143     int diff;
00144     int smallest = INT_MAX;
00145     int tmp, tmp2;
00146 
00147     #define CMP(i, j) deshake->c.sad[0](deshake, src1 + cy * stride + cx, \
00148                                         src2 + (j) * stride + (i), stride, \
00149                                         deshake->blocksize)
00150 
00151     if (deshake->search == EXHAUSTIVE) {
00152         // Compare every possible position - this is sloooow!
00153         for (y = -deshake->ry; y <= deshake->ry; y++) {
00154             for (x = -deshake->rx; x <= deshake->rx; x++) {
00155                 diff = CMP(cx - x, cy - y);
00156                 if (diff < smallest) {
00157                     smallest = diff;
00158                     mv->x = x;
00159                     mv->y = y;
00160                 }
00161             }
00162         }
00163     } else if (deshake->search == SMART_EXHAUSTIVE) {
00164         // Compare every other possible position and find the best match
00165         for (y = -deshake->ry + 1; y < deshake->ry - 2; y += 2) {
00166             for (x = -deshake->rx + 1; x < deshake->rx - 2; x += 2) {
00167                 diff = CMP(cx - x, cy - y);
00168                 if (diff < smallest) {
00169                     smallest = diff;
00170                     mv->x = x;
00171                     mv->y = y;
00172                 }
00173             }
00174         }
00175 
00176         // Hone in on the specific best match around the match we found above
00177         tmp = mv->x;
00178         tmp2 = mv->y;
00179 
00180         for (y = tmp2 - 1; y <= tmp2 + 1; y++) {
00181             for (x = tmp - 1; x <= tmp + 1; x++) {
00182                 if (x == tmp && y == tmp2)
00183                     continue;
00184 
00185                 diff = CMP(cx - x, cy - y);
00186                 if (diff < smallest) {
00187                     smallest = diff;
00188                     mv->x = x;
00189                     mv->y = y;
00190                 }
00191             }
00192         }
00193     }
00194 
00195     if (smallest > 512) {
00196         mv->x = -1;
00197         mv->y = -1;
00198     }
00199     emms_c();
00200     //av_log(NULL, AV_LOG_ERROR, "%d\n", smallest);
00201     //av_log(NULL, AV_LOG_ERROR, "Final: (%d, %d) = %d x %d\n", cx, cy, mv->x, mv->y);
00202 }
00203 
00209 static int block_contrast(uint8_t *src, int x, int y, int stride, int blocksize)
00210 {
00211     int highest = 0;
00212     int lowest = 0;
00213     int i, j, pos;
00214 
00215     for (i = 0; i <= blocksize * 2; i++) {
00216         // We use a width of 16 here to match the libavcodec sad functions
00217         for (j = 0; i <= 15; i++) {
00218             pos = (y - i) * stride + (x - j);
00219             if (src[pos] < lowest)
00220                 lowest = src[pos];
00221             else if (src[pos] > highest) {
00222                 highest = src[pos];
00223             }
00224         }
00225     }
00226 
00227     return highest - lowest;
00228 }
00229 
00233 static double block_angle(int x, int y, int cx, int cy, IntMotionVector *shift)
00234 {
00235     double a1, a2, diff;
00236 
00237     a1 = atan2(y - cy, x - cx);
00238     a2 = atan2(y - cy + shift->y, x - cx + shift->x);
00239 
00240     diff = a2 - a1;
00241 
00242     return (diff > M_PI)  ? diff - 2 * M_PI :
00243            (diff < -M_PI) ? diff + 2 * M_PI :
00244            diff;
00245 }
00246 
00254 static void find_motion(DeshakeContext *deshake, uint8_t *src1, uint8_t *src2,
00255                         int width, int height, int stride, Transform *t)
00256 {
00257     int x, y;
00258     IntMotionVector mv = {0, 0};
00259     int counts[128][128];
00260     int count_max_value = 0;
00261     int contrast;
00262 
00263     int pos;
00264     double *angles = av_malloc(sizeof(*angles) * width * height / (16 * deshake->blocksize));
00265     int center_x = 0, center_y = 0;
00266     double p_x, p_y;
00267 
00268     // Reset counts to zero
00269     for (x = 0; x < deshake->rx * 2 + 1; x++) {
00270         for (y = 0; y < deshake->ry * 2 + 1; y++) {
00271             counts[x][y] = 0;
00272         }
00273     }
00274 
00275     pos = 0;
00276     // Find motion for every block and store the motion vector in the counts
00277     for (y = deshake->ry; y < height - deshake->ry - (deshake->blocksize * 2); y += deshake->blocksize * 2) {
00278         // We use a width of 16 here to match the libavcodec sad functions
00279         for (x = deshake->rx; x < width - deshake->rx - 16; x += 16) {
00280             // If the contrast is too low, just skip this block as it probably
00281             // won't be very useful to us.
00282             contrast = block_contrast(src2, x, y, stride, deshake->blocksize);
00283             if (contrast > deshake->contrast) {
00284                 //av_log(NULL, AV_LOG_ERROR, "%d\n", contrast);
00285                 find_block_motion(deshake, src1, src2, x, y, stride, &mv);
00286                 if (mv.x != -1 && mv.y != -1) {
00287                     counts[mv.x + deshake->rx][mv.y + deshake->ry] += 1;
00288                     if (x > deshake->rx && y > deshake->ry)
00289                         angles[pos++] = block_angle(x, y, 0, 0, &mv);
00290 
00291                     center_x += mv.x;
00292                     center_y += mv.y;
00293                 }
00294             }
00295         }
00296     }
00297 
00298     if (pos) {
00299          center_x /= pos;
00300          center_y /= pos;
00301          t->angle = clean_mean(angles, pos);
00302          if (t->angle < 0.001)
00303               t->angle = 0;
00304     } else {
00305          t->angle = 0;
00306     }
00307 
00308     // Find the most common motion vector in the frame and use it as the gmv
00309     for (y = deshake->ry * 2; y >= 0; y--) {
00310         for (x = 0; x < deshake->rx * 2 + 1; x++) {
00311             //av_log(NULL, AV_LOG_ERROR, "%5d ", counts[x][y]);
00312             if (counts[x][y] > count_max_value) {
00313                 t->vector.x = x - deshake->rx;
00314                 t->vector.y = y - deshake->ry;
00315                 count_max_value = counts[x][y];
00316             }
00317         }
00318         //av_log(NULL, AV_LOG_ERROR, "\n");
00319     }
00320 
00321     p_x = (center_x - width / 2);
00322     p_y = (center_y - height / 2);
00323     t->vector.x += (cos(t->angle)-1)*p_x  - sin(t->angle)*p_y;
00324     t->vector.y += sin(t->angle)*p_x  + (cos(t->angle)-1)*p_y;
00325 
00326     // Clamp max shift & rotation?
00327     t->vector.x = av_clipf(t->vector.x, -deshake->rx * 2, deshake->rx * 2);
00328     t->vector.y = av_clipf(t->vector.y, -deshake->ry * 2, deshake->ry * 2);
00329     t->angle = av_clipf(t->angle, -0.1, 0.1);
00330 
00331     //av_log(NULL, AV_LOG_ERROR, "%d x %d\n", avg->x, avg->y);
00332     av_free(angles);
00333 }
00334 
00335 static av_cold int init(AVFilterContext *ctx, const char *args)
00336 {
00337     DeshakeContext *deshake = ctx->priv;
00338     char filename[256] = {0};
00339 
00340     deshake->rx = 16;
00341     deshake->ry = 16;
00342     deshake->edge = FILL_MIRROR;
00343     deshake->blocksize = 8;
00344     deshake->contrast = 125;
00345     deshake->search = EXHAUSTIVE;
00346     deshake->refcount = 20;
00347 
00348     deshake->cw = -1;
00349     deshake->ch = -1;
00350     deshake->cx = -1;
00351     deshake->cy = -1;
00352 
00353     if (args) {
00354         sscanf(args, "%d:%d:%d:%d:%d:%d:%d:%d:%d:%d:%255s",
00355                &deshake->cx, &deshake->cy, &deshake->cw, &deshake->ch,
00356                &deshake->rx, &deshake->ry, (int *)&deshake->edge,
00357                &deshake->blocksize, &deshake->contrast, (int *)&deshake->search, filename);
00358 
00359         deshake->blocksize /= 2;
00360 
00361         deshake->rx = av_clip(deshake->rx, 0, 64);
00362         deshake->ry = av_clip(deshake->ry, 0, 64);
00363         deshake->edge = av_clip(deshake->edge, FILL_BLANK, FILL_COUNT - 1);
00364         deshake->blocksize = av_clip(deshake->blocksize, 4, 128);
00365         deshake->contrast = av_clip(deshake->contrast, 1, 255);
00366         deshake->search = av_clip(deshake->search, EXHAUSTIVE, SEARCH_COUNT - 1);
00367 
00368     }
00369     if (*filename)
00370         deshake->fp = fopen(filename, "w");
00371     if (deshake->fp)
00372         fwrite("Ori x, Avg x, Fin x, Ori y, Avg y, Fin y, Ori angle, Avg angle, Fin angle, Ori zoom, Avg zoom, Fin zoom\n", sizeof(char), 104, deshake->fp);
00373 
00374     // Quadword align left edge of box for MMX code, adjust width if necessary
00375     // to keep right margin
00376     if (deshake->cx > 0) {
00377         deshake->cw += deshake->cx - (deshake->cx & ~15);
00378         deshake->cx &= ~15;
00379     }
00380 
00381     av_log(ctx, AV_LOG_VERBOSE, "cx: %d, cy: %d, cw: %d, ch: %d, rx: %d, ry: %d, edge: %d blocksize: %d contrast: %d search: %d\n",
00382            deshake->cx, deshake->cy, deshake->cw, deshake->ch,
00383            deshake->rx, deshake->ry, deshake->edge, deshake->blocksize * 2, deshake->contrast, deshake->search);
00384 
00385     return 0;
00386 }
00387 
00388 static int query_formats(AVFilterContext *ctx)
00389 {
00390     enum AVPixelFormat pix_fmts[] = {
00391         AV_PIX_FMT_YUV420P,  AV_PIX_FMT_YUV422P,  AV_PIX_FMT_YUV444P,  AV_PIX_FMT_YUV410P,
00392         AV_PIX_FMT_YUV411P,  AV_PIX_FMT_YUV440P,  AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ422P,
00393         AV_PIX_FMT_YUVJ444P, AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_NONE
00394     };
00395 
00396     ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
00397 
00398     return 0;
00399 }
00400 
00401 static int config_props(AVFilterLink *link)
00402 {
00403     DeshakeContext *deshake = link->dst->priv;
00404 
00405     deshake->ref = NULL;
00406     deshake->last.vector.x = 0;
00407     deshake->last.vector.y = 0;
00408     deshake->last.angle = 0;
00409     deshake->last.zoom = 0;
00410 
00411     deshake->avctx = avcodec_alloc_context3(NULL);
00412     dsputil_init(&deshake->c, deshake->avctx);
00413 
00414     return 0;
00415 }
00416 
00417 static av_cold void uninit(AVFilterContext *ctx)
00418 {
00419     DeshakeContext *deshake = ctx->priv;
00420 
00421     avfilter_unref_buffer(deshake->ref);
00422     if (deshake->fp)
00423         fclose(deshake->fp);
00424     if (deshake->avctx)
00425         avcodec_close(deshake->avctx);
00426     av_freep(&deshake->avctx);
00427 }
00428 
00429 static int filter_frame(AVFilterLink *link, AVFilterBufferRef *in)
00430 {
00431     DeshakeContext *deshake = link->dst->priv;
00432     AVFilterLink *outlink = link->dst->outputs[0];
00433     AVFilterBufferRef *out;
00434     Transform t = {{0},0}, orig = {{0},0};
00435     float matrix[9];
00436     float alpha = 2.0 / deshake->refcount;
00437     char tmp[256];
00438 
00439     out = ff_get_video_buffer(outlink, AV_PERM_WRITE, outlink->w, outlink->h);
00440     if (!out) {
00441         avfilter_unref_bufferp(&in);
00442         return AVERROR(ENOMEM);
00443     }
00444     avfilter_copy_buffer_ref_props(out, in);
00445 
00446     link->cur_buf = NULL; /* it is in 'in' now */
00447     if (deshake->cx < 0 || deshake->cy < 0 || deshake->cw < 0 || deshake->ch < 0) {
00448         // Find the most likely global motion for the current frame
00449         find_motion(deshake, (deshake->ref == NULL) ? in->data[0] : deshake->ref->data[0], in->data[0], link->w, link->h, in->linesize[0], &t);
00450     } else {
00451         uint8_t *src1 = (deshake->ref == NULL) ? in->data[0] : deshake->ref->data[0];
00452         uint8_t *src2 = in->data[0];
00453 
00454         deshake->cx = FFMIN(deshake->cx, link->w);
00455         deshake->cy = FFMIN(deshake->cy, link->h);
00456 
00457         if ((unsigned)deshake->cx + (unsigned)deshake->cw > link->w) deshake->cw = link->w - deshake->cx;
00458         if ((unsigned)deshake->cy + (unsigned)deshake->ch > link->h) deshake->ch = link->h - deshake->cy;
00459 
00460         // Quadword align right margin
00461         deshake->cw &= ~15;
00462 
00463         src1 += deshake->cy * in->linesize[0] + deshake->cx;
00464         src2 += deshake->cy * in->linesize[0] + deshake->cx;
00465 
00466         find_motion(deshake, src1, src2, deshake->cw, deshake->ch, in->linesize[0], &t);
00467     }
00468 
00469 
00470     // Copy transform so we can output it later to compare to the smoothed value
00471     orig.vector.x = t.vector.x;
00472     orig.vector.y = t.vector.y;
00473     orig.angle = t.angle;
00474     orig.zoom = t.zoom;
00475 
00476     // Generate a one-sided moving exponential average
00477     deshake->avg.vector.x = alpha * t.vector.x + (1.0 - alpha) * deshake->avg.vector.x;
00478     deshake->avg.vector.y = alpha * t.vector.y + (1.0 - alpha) * deshake->avg.vector.y;
00479     deshake->avg.angle = alpha * t.angle + (1.0 - alpha) * deshake->avg.angle;
00480     deshake->avg.zoom = alpha * t.zoom + (1.0 - alpha) * deshake->avg.zoom;
00481 
00482     // Remove the average from the current motion to detect the motion that
00483     // is not on purpose, just as jitter from bumping the camera
00484     t.vector.x -= deshake->avg.vector.x;
00485     t.vector.y -= deshake->avg.vector.y;
00486     t.angle -= deshake->avg.angle;
00487     t.zoom -= deshake->avg.zoom;
00488 
00489     // Invert the motion to undo it
00490     t.vector.x *= -1;
00491     t.vector.y *= -1;
00492     t.angle *= -1;
00493 
00494     // Write statistics to file
00495     if (deshake->fp) {
00496         snprintf(tmp, 256, "%f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f, %f\n", orig.vector.x, deshake->avg.vector.x, t.vector.x, orig.vector.y, deshake->avg.vector.y, t.vector.y, orig.angle, deshake->avg.angle, t.angle, orig.zoom, deshake->avg.zoom, t.zoom);
00497         fwrite(tmp, sizeof(char), strlen(tmp), deshake->fp);
00498     }
00499 
00500     // Turn relative current frame motion into absolute by adding it to the
00501     // last absolute motion
00502     t.vector.x += deshake->last.vector.x;
00503     t.vector.y += deshake->last.vector.y;
00504     t.angle += deshake->last.angle;
00505     t.zoom += deshake->last.zoom;
00506 
00507     // Shrink motion by 10% to keep things centered in the camera frame
00508     t.vector.x *= 0.9;
00509     t.vector.y *= 0.9;
00510     t.angle *= 0.9;
00511 
00512     // Store the last absolute motion information
00513     deshake->last.vector.x = t.vector.x;
00514     deshake->last.vector.y = t.vector.y;
00515     deshake->last.angle = t.angle;
00516     deshake->last.zoom = t.zoom;
00517 
00518     // Generate a luma transformation matrix
00519     avfilter_get_matrix(t.vector.x, t.vector.y, t.angle, 1.0 + t.zoom / 100.0, matrix);
00520 
00521     // Transform the luma plane
00522     avfilter_transform(in->data[0], out->data[0], in->linesize[0], out->linesize[0], link->w, link->h, matrix, INTERPOLATE_BILINEAR, deshake->edge);
00523 
00524     // Generate a chroma transformation matrix
00525     avfilter_get_matrix(t.vector.x / (link->w / CHROMA_WIDTH(link)), t.vector.y / (link->h / CHROMA_HEIGHT(link)), t.angle, 1.0 + t.zoom / 100.0, matrix);
00526 
00527     // Transform the chroma planes
00528     avfilter_transform(in->data[1], out->data[1], in->linesize[1], out->linesize[1], CHROMA_WIDTH(link), CHROMA_HEIGHT(link), matrix, INTERPOLATE_BILINEAR, deshake->edge);
00529     avfilter_transform(in->data[2], out->data[2], in->linesize[2], out->linesize[2], CHROMA_WIDTH(link), CHROMA_HEIGHT(link), matrix, INTERPOLATE_BILINEAR, deshake->edge);
00530 
00531     // Cleanup the old reference frame
00532     avfilter_unref_buffer(deshake->ref);
00533 
00534     // Store the current frame as the reference frame for calculating the
00535     // motion of the next frame
00536     deshake->ref = in;
00537 
00538     return ff_filter_frame(outlink, out);
00539 }
00540 
00541 static const AVFilterPad deshake_inputs[] = {
00542     {
00543         .name         = "default",
00544         .type         = AVMEDIA_TYPE_VIDEO,
00545         .filter_frame = filter_frame,
00546         .config_props = config_props,
00547         .min_perms    = AV_PERM_READ | AV_PERM_PRESERVE,
00548     },
00549     { NULL }
00550 };
00551 
00552 static const AVFilterPad deshake_outputs[] = {
00553     {
00554         .name = "default",
00555         .type = AVMEDIA_TYPE_VIDEO,
00556     },
00557     { NULL }
00558 };
00559 
00560 AVFilter avfilter_vf_deshake = {
00561     .name          = "deshake",
00562     .description   = NULL_IF_CONFIG_SMALL("Stabilize shaky video."),
00563     .priv_size     = sizeof(DeshakeContext),
00564     .init          = init,
00565     .uninit        = uninit,
00566     .query_formats = query_formats,
00567     .inputs        = deshake_inputs,
00568     .outputs       = deshake_outputs,
00569 };