git.sesse.net Git - ffmpeg/blob - libavformat/chromaprint.c

   1 /*
   2  * Chromaprint fingerprinting muxer
   3  * Copyright (c) 2015 Rodger Combs
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 #include "avformat.h"
  23 #include "libavutil/opt.h"
  24 #include "libavcodec/internal.h"
  25 #include <chromaprint.h>
  26
  27 #define CPR_VERSION_INT AV_VERSION_INT(CHROMAPRINT_VERSION_MAJOR, \
  28                                        CHROMAPRINT_VERSION_MINOR, \
  29                                        CHROMAPRINT_VERSION_PATCH)
  30
  31 typedef enum FingerprintFormat {
  32     FINGERPRINT_RAW,
  33     FINGERPRINT_COMPRESSED,
  34     FINGERPRINT_BASE64,
  35 } FingerprintFormat;
  36
  37 typedef struct ChromaprintMuxContext {
  38     const AVClass *class;
  39     int silence_threshold;
  40     int algorithm;
  41     FingerprintFormat fp_format;
  42     ChromaprintContext ctx;
  43 } ChromaprintMuxContext;
  44
  45 static void cleanup(ChromaprintMuxContext *cpr)
  46 {
  47     if (cpr->ctx) {
  48         avpriv_lock_avformat();
  49         chromaprint_free(cpr->ctx);
  50         avpriv_unlock_avformat();
  51     }
  52 }
  53
  54 static int write_header(AVFormatContext *s)
  55 {
  56     ChromaprintMuxContext *cpr = s->priv_data;
  57     AVStream *st;
  58
  59     avpriv_lock_avformat();
  60     cpr->ctx = chromaprint_new(cpr->algorithm);
  61     avpriv_unlock_avformat();
  62
  63     if (!cpr->ctx) {
  64         av_log(s, AV_LOG_ERROR, "Failed to create chromaprint context.\n");
  65         return AVERROR(ENOMEM);
  66     }
  67
  68     if (cpr->silence_threshold != -1) {
  69 #if CPR_VERSION_INT >= AV_VERSION_INT(0, 7, 0)
  70         if (!chromaprint_set_option(cpr->ctx, "silence_threshold", cpr->silence_threshold)) {
  71             av_log(s, AV_LOG_ERROR, "Failed to set silence threshold.\n");
  72             goto fail;
  73         }
  74 #else
  75         av_log(s, AV_LOG_ERROR, "Setting the silence threshold requires Chromaprint "
  76                                 "version 0.7.0 or later.\n");
  77         goto fail;
  78 #endif
  79     }
  80
  81     if (s->nb_streams != 1) {
  82         av_log(s, AV_LOG_ERROR, "Only one stream is supported\n");
  83         goto fail;
  84     }
  85
  86     st = s->streams[0];
  87
  88     if (st->codec->channels > 2) {
  89         av_log(s, AV_LOG_ERROR, "Only up to 2 channels are supported\n");
  90         goto fail;
  91     }
  92
  93     if (st->codec->sample_rate < 1000) {
  94         av_log(s, AV_LOG_ERROR, "Sampling rate must be at least 1000\n");
  95         goto fail;
  96     }
  97
  98     if (!chromaprint_start(cpr->ctx, st->codec->sample_rate, st->codec->channels)) {
  99         av_log(s, AV_LOG_ERROR, "Failed to start chromaprint\n");
 100         goto fail;
 101     }
 102
 103     return 0;
 104 fail:
 105     cleanup(cpr);
 106     return AVERROR(EINVAL);
 107 }
 108
 109 static int write_packet(AVFormatContext *s, AVPacket *pkt)
 110 {
 111     ChromaprintMuxContext *cpr = s->priv_data;
 112     return chromaprint_feed(cpr->ctx, pkt->data, pkt->size / 2) ? 0 : AVERROR(EINVAL);
 113 }
 114
 115 static int write_trailer(AVFormatContext *s)
 116 {
 117     ChromaprintMuxContext *cpr = s->priv_data;
 118     AVIOContext *pb = s->pb;
 119     void *fp = NULL, *enc_fp = NULL;
 120     int size, enc_size, ret = AVERROR(EINVAL);
 121
 122     if (!chromaprint_finish(cpr->ctx)) {
 123         av_log(s, AV_LOG_ERROR, "Failed to generate fingerprint\n");
 124         goto fail;
 125     }
 126
 127     if (!chromaprint_get_raw_fingerprint(cpr->ctx, &fp, &size)) {
 128         av_log(s, AV_LOG_ERROR, "Failed to retrieve fingerprint\n");
 129         goto fail;
 130     }
 131
 132     switch (cpr->fp_format) {
 133     case FINGERPRINT_RAW:
 134         avio_write(pb, fp, size);
 135         break;
 136     case FINGERPRINT_COMPRESSED:
 137     case FINGERPRINT_BASE64:
 138         if (!chromaprint_encode_fingerprint(fp, size, cpr->algorithm, &enc_fp, &enc_size,
 139                                             cpr->fp_format == FINGERPRINT_BASE64)) {
 140             av_log(s, AV_LOG_ERROR, "Failed to encode fingerprint\n");
 141             goto fail;
 142         }
 143         avio_write(pb, enc_fp, enc_size);
 144         break;
 145     }
 146
 147     ret = 0;
 148 fail:
 149     if (fp)
 150         chromaprint_dealloc(fp);
 151     if (enc_fp)
 152         chromaprint_dealloc(enc_fp);
 153     cleanup(cpr);
 154     return ret;
 155 }
 156
 157 #define OFFSET(x) offsetof(ChromaprintMuxContext, x)
 158 #define FLAGS AV_OPT_FLAG_ENCODING_PARAM
 159 static const AVOption options[] = {
 160     { "silence_threshold", "threshold for detecting silence", OFFSET(silence_threshold), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 32767, FLAGS },
 161     { "algorithm", "version of the fingerprint algorithm", OFFSET(algorithm), AV_OPT_TYPE_INT, { .i64 = CHROMAPRINT_ALGORITHM_DEFAULT }, CHROMAPRINT_ALGORITHM_TEST1, INT_MAX, FLAGS },
 162     { "fp_format", "fingerprint format to write", OFFSET(fp_format), AV_OPT_TYPE_INT, { .i64 = FINGERPRINT_BASE64 }, FINGERPRINT_RAW, FINGERPRINT_BASE64, FLAGS },
 163     { "raw", "binary raw fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_RAW }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
 164     { "compressed", "binary compressed fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_COMPRESSED }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
 165     { "base64", "Base64 compressed fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_BASE64 }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
 166     { NULL },
 167 };
 168
 169 static const AVClass chromaprint_class = {
 170     .class_name = "chromaprint muxer",
 171     .item_name  = av_default_item_name,
 172     .option     = options,
 173     .version    = LIBAVUTIL_VERSION_INT,
 174 };
 175
 176 AVOutputFormat ff_chromaprint_muxer = {
 177     .name              = "chromaprint",
 178     .long_name         = NULL_IF_CONFIG_SMALL("Chromaprint"),
 179     .priv_data_size    = sizeof(ChromaprintMuxContext),
 180     .audio_codec       = AV_NE(AV_CODEC_ID_PCM_S16BE, AV_CODEC_ID_PCM_S16LE),
 181     .write_header      = write_header,
 182     .write_packet      = write_packet,
 183     .write_trailer     = write_trailer,
 184     .flags             = AVFMT_NOTIMESTAMPS,
 185     .priv_class        = &chromaprint_class,
 186 };