1/*
2 * Chromaprint fingerprinting muxer
3 * Copyright (c) 2015 rcombs
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 "internal.h"
24#include "libavutil/opt.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
31typedef enum FingerprintFormat {
32    FINGERPRINT_RAW,
33    FINGERPRINT_COMPRESSED,
34    FINGERPRINT_BASE64,
35} FingerprintFormat;
36
37typedef struct ChromaprintMuxContext {
38    const AVClass *class;
39    int silence_threshold;
40    int algorithm;
41    FingerprintFormat fp_format;
42#if CPR_VERSION_INT >= AV_VERSION_INT(1, 4, 0)
43    ChromaprintContext *ctx;
44#else
45    ChromaprintContext ctx;
46#endif
47} ChromaprintMuxContext;
48
49static void deinit(AVFormatContext *s)
50{
51    ChromaprintMuxContext *const cpr = s->priv_data;
52
53    if (cpr->ctx) {
54        ff_lock_avformat();
55        chromaprint_free(cpr->ctx);
56        ff_unlock_avformat();
57    }
58}
59
60static int write_header(AVFormatContext *s)
61{
62    ChromaprintMuxContext *cpr = s->priv_data;
63    AVStream *st;
64
65    ff_lock_avformat();
66    cpr->ctx = chromaprint_new(cpr->algorithm);
67    ff_unlock_avformat();
68
69    if (!cpr->ctx) {
70        av_log(s, AV_LOG_ERROR, "Failed to create chromaprint context.\n");
71        return AVERROR_EXTERNAL;
72    }
73
74    if (cpr->silence_threshold != -1) {
75#if CPR_VERSION_INT >= AV_VERSION_INT(0, 7, 0)
76        if (!chromaprint_set_option(cpr->ctx, "silence_threshold", cpr->silence_threshold)) {
77            av_log(s, AV_LOG_ERROR, "Failed to set silence threshold. Setting silence_threshold requires -algorithm 3 option.\n");
78            return AVERROR_EXTERNAL;
79        }
80#else
81        av_log(s, AV_LOG_ERROR, "Setting the silence threshold requires Chromaprint "
82                                "version 0.7.0 or later.\n");
83        return AVERROR(ENOSYS);
84#endif
85    }
86
87    if (s->nb_streams != 1) {
88        av_log(s, AV_LOG_ERROR, "Only one stream is supported\n");
89        return AVERROR(EINVAL);
90    }
91
92    st = s->streams[0];
93
94    if (st->codecpar->ch_layout.nb_channels > 2) {
95        av_log(s, AV_LOG_ERROR, "Only up to 2 channels are supported\n");
96        return AVERROR(EINVAL);
97    }
98
99    if (st->codecpar->sample_rate < 1000) {
100        av_log(s, AV_LOG_ERROR, "Sampling rate must be at least 1000\n");
101        return AVERROR(EINVAL);
102    }
103
104    if (!chromaprint_start(cpr->ctx, st->codecpar->sample_rate, st->codecpar->ch_layout.nb_channels)) {
105        av_log(s, AV_LOG_ERROR, "Failed to start chromaprint\n");
106        return AVERROR_EXTERNAL;
107    }
108
109    return 0;
110}
111
112static int write_packet(AVFormatContext *s, AVPacket *pkt)
113{
114    ChromaprintMuxContext *cpr = s->priv_data;
115    return chromaprint_feed(cpr->ctx, (const int16_t *)pkt->data, pkt->size / 2) ? 0 : AVERROR(EINVAL);
116}
117
118static int write_trailer(AVFormatContext *s)
119{
120    ChromaprintMuxContext *cpr = s->priv_data;
121    AVIOContext *pb = s->pb;
122    void *fp = NULL;
123    char *enc_fp = NULL;
124    int size, enc_size, ret = AVERROR_EXTERNAL;
125
126    if (!chromaprint_finish(cpr->ctx)) {
127        av_log(s, AV_LOG_ERROR, "Failed to generate fingerprint\n");
128        goto fail;
129    }
130
131    if (!chromaprint_get_raw_fingerprint(cpr->ctx, (uint32_t **)&fp, &size)) {
132        av_log(s, AV_LOG_ERROR, "Failed to retrieve fingerprint\n");
133        goto fail;
134    }
135
136    switch (cpr->fp_format) {
137    case FINGERPRINT_RAW:
138        avio_write(pb, fp, size * 4); //fp points to array of uint32_t
139        break;
140    case FINGERPRINT_COMPRESSED:
141    case FINGERPRINT_BASE64:
142        if (!chromaprint_encode_fingerprint(fp, size, cpr->algorithm, &enc_fp, &enc_size,
143                                            cpr->fp_format == FINGERPRINT_BASE64)) {
144            av_log(s, AV_LOG_ERROR, "Failed to encode fingerprint\n");
145            goto fail;
146        }
147        avio_write(pb, enc_fp, enc_size);
148        break;
149    }
150
151    ret = 0;
152fail:
153    if (fp)
154        chromaprint_dealloc(fp);
155    if (enc_fp)
156        chromaprint_dealloc(enc_fp);
157    return ret;
158}
159
160#define OFFSET(x) offsetof(ChromaprintMuxContext, x)
161#define FLAGS AV_OPT_FLAG_ENCODING_PARAM
162static const AVOption options[] = {
163    { "silence_threshold", "threshold for detecting silence", OFFSET(silence_threshold), AV_OPT_TYPE_INT, { .i64 = -1 }, -1, 32767, FLAGS },
164    { "algorithm", "version of the fingerprint algorithm", OFFSET(algorithm), AV_OPT_TYPE_INT, { .i64 = CHROMAPRINT_ALGORITHM_DEFAULT }, CHROMAPRINT_ALGORITHM_TEST1, INT_MAX, FLAGS },
165    { "fp_format", "fingerprint format to write", OFFSET(fp_format), AV_OPT_TYPE_INT, { .i64 = FINGERPRINT_BASE64 }, FINGERPRINT_RAW, FINGERPRINT_BASE64, FLAGS, "fp_format" },
166    { "raw", "binary raw fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_RAW }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
167    { "compressed", "binary compressed fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_COMPRESSED }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
168    { "base64", "Base64 compressed fingerprint", 0, AV_OPT_TYPE_CONST, {.i64 = FINGERPRINT_BASE64 }, INT_MIN, INT_MAX, FLAGS, "fp_format"},
169    { NULL },
170};
171
172static const AVClass chromaprint_class = {
173    .class_name = "chromaprint muxer",
174    .item_name  = av_default_item_name,
175    .option     = options,
176    .version    = LIBAVUTIL_VERSION_INT,
177};
178
179const AVOutputFormat ff_chromaprint_muxer = {
180    .name              = "chromaprint",
181    .long_name         = NULL_IF_CONFIG_SMALL("Chromaprint"),
182    .priv_data_size    = sizeof(ChromaprintMuxContext),
183    .audio_codec       = AV_NE(AV_CODEC_ID_PCM_S16BE, AV_CODEC_ID_PCM_S16LE),
184    .write_header      = write_header,
185    .write_packet      = write_packet,
186    .write_trailer     = write_trailer,
187    .deinit            = deinit,
188    .flags             = AVFMT_NOTIMESTAMPS,
189    .priv_class        = &chromaprint_class,
190};
191