view utilities/track/track.cpp @ 33:84d6acb6b3ba

* Bit more work on track composer identification

/* -*- c-basic-offset: 4 indent-tabs-mode: nil -*-  vi:set ts=8 sts=4 sw=4: */

#include "Objects.h"
#include "TypeRegistrar.h"

#include <dataquay/BasicStore.h>
#include <dataquay/TransactionalStore.h>
#include <dataquay/RDFException.h>
#include <dataquay/objectmapper/ObjectLoader.h>
#include <dataquay/objectmapper/ObjectStorer.h>
#include <dataquay/objectmapper/ObjectMapper.h>
#include <dataquay/objectmapper/TypeMapping.h>
#include <dataquay/Debug.h>

#include "data/fileio/AudioFileReaderFactory.h"
#include "data/fileio/AudioFileReader.h"
#include "base/TempDirectory.h"

#include "Matcher.h"

#include <vamp-hostsdk/PluginLoader.h>

#include <QMultiMap>
#include <QFileInfo>

#include <iostream>

using namespace Dataquay;
using namespace ClassicalData;
using namespace std;
using namespace Vamp;
using namespace Vamp::HostExt;

ostream &operator<<(ostream &target, const QString &str)
{
    return target << str.toLocal8Bit().data();
}

ostream &operator<<(ostream &target, const QUrl &u)
{
    return target << "<" << u.toString() << ">";
}


bool
load(BasicStore *store, QString fileName)
{
    QUrl url = QUrl::fromLocalFile(fileName);

    cerr << "Importing from URL " << url << " ...";
    try {
	store->import(url, BasicStore::ImportPermitDuplicates);
    } catch (RDFException e) {
        cerr << " retrying with explicit ntriples type...";
        try {
            store->import(url, BasicStore::ImportPermitDuplicates, "ntriples");
        } catch (RDFException e) {
            cerr << "failed" << endl;
            cerr << "Import failed: " << e.what() << endl;
            return false;
        }
    }

    cerr << " done" << endl;
    return true;
}

void
usage(char *name)
{
    int s = 0;
    for (int i = 0; name[i]; ++i) if (name[i] == '/') s = i + 1;
    name = name + s;
    cerr << "Usage:" << endl;
    cerr << "  " << name << " <input-rdf-file> guess <track.wav> [<track.wav> ...]" << endl;
    exit(-1);
}

static QList<Composer *> allComposers;
static QHash<QString, Composer *> composerAliases;
static QHash<Uri, Composer *> composerUris;
static QMap<Composer *, QSet<Work *> > worksMap;

void
show(Composer *c)
{
    cout << c->property("uri").value<Uri>() << endl;
    cout << c->getSortName(true);
    QString d = c->getDisplayDates();
    if (d != "") cout << " (" << d << ")";
    if (!c->nationality().empty() || c->period() != "") {
        cout << " [";
        bool first = true;
        foreach (QString n, c->nationality()) {
            if (!first) cout << "/";
            cout << n;
            first = false;
        }
        if (c->period() != "") {
            if (!first) cout << ", ";
            cout << c->period();
        }
        cout << "]";
    }
    if (c->gender() != "") {
        cout << " *" << c->gender();
    }
    if (!worksMap[c].empty()) {
        cout << " [" << worksMap[c].size() << " work(s)]";
    }
    cout << endl;
    foreach (QString a, c->aliases()) {
        cout << " - " << a << endl;
    }
    if (c->remarks() != "") {
        cout << " " << c->remarks() << endl;
    }
    foreach (Document *d, c->pages()) {
        cout << d->siteName() << " -> " << d->uri() << endl;
    }
    foreach (Uri u, c->otherURIs()) {
        cout << "Same as " << u << endl;
    }
}

void
showBrief(Composer *c)
{
    cout << c->property("uri").value<Uri>() << endl;
    cout << c->getSortName(false);
    QString d = c->getDisplayDates();
    if (d != "") cout << " (" << d << ")";
    if (!c->nationality().empty() || c->period() != "") {
        cout << " [";
        bool first = true;
        foreach (QString n, c->nationality()) {
            if (!first) cout << "/";
            cout << n;
            first = false;
        }
        if (c->period() != "") {
            if (!first) cout << " ";
            cout << c->period();
        }
        cout << "]";
    }
    if (c->gender() != "") {
        cout << " *" << c->gender();
    }
    if (!worksMap[c].empty()) {
        cout << " [" << worksMap[c].size() << " work(s)]";
    }
    cout << endl;
}

void
listBrief(QList<Composer *> composers)
{
    QMultiMap<QString, Composer *> sorted;
    foreach (Composer *c, composers) {
        sorted.insert(c->getSortName(false), c);
    }
    foreach (Composer *c, sorted) {
        showBrief(c);
    }
}

void
listUris(QList<Composer *> composers)
{
    QMultiMap<Uri, Composer *> sorted;
    foreach (Composer *c, composers) {
        sorted.insert(c->property("uri").value<Uri>(), c);
    }
    foreach (Uri uri, sorted.keys()) {
        cout << uri << endl;
    }
}

void
showSearchResults(QMultiMap<float, Composer *> matches, int count)
{
    int n = 0;
    for (QMultiMap<float, Composer *>::const_iterator i = matches.end();
         i != matches.begin(); ) {
        --i;
        if (i.key() <= 0) continue;
        cout << endl;
        if (n == 0) {
            cout << "Best match:" << endl;
        } else if (n == 1) {
            cout << "Other candidate(s):" << endl;
        }
        cout << "[" << i.key() << "] ";
        if (n == 0) show(i.value());
        else showBrief(i.value());
        if (++n > count) break;
    }
    if (n == 0) cout << "No matches" << endl;
    cout << endl;
}

void
getTrackData(FileSource source, QString &fingerprint, QString &puid,
             QString &title, QString &maker, AudioFileReader::TagMap &tags)
{
    AudioFileReader *reader = AudioFileReaderFactory::createReader(source);
    if (!reader || !reader->isOK()) {
        cerr << "Failed to open audio file" << endl;
        return;
    }

    title = reader->getTitle();
    maker = reader->getMaker();
    cout << "File tag title: " << reader->getTitle() << endl;
    cout << "File tag maker: " << reader->getMaker() << endl;

    cout << "All tags:" << endl;
    tags = reader->getTags();
    for (AudioFileReader::TagMap::const_iterator i = tags.begin();
         i != tags.end(); ++i) {
        cout << i->first << " " << i->second << endl;
    }

    PluginLoader *pl = PluginLoader::getInstance();
    Plugin *plugin = pl->loadPlugin
        ("ofa-vamp-plugin:ofa_puid_and_fingerprint", reader->getSampleRate(), PluginLoader::ADAPT_ALL);
    if (!plugin) {
        cerr << "Failed to load OFA Vamp plugin" << endl;
        return;
    }

    // 135 seconds... well, ok, let's have 136
    int secs = 136;
    
    int want = int(secs * reader->getSampleRate());
    int ch = reader->getChannelCount();
    std::vector<SampleBlock> samples;
    reader->getDeInterleavedFrames(0, want, samples);
    int have = samples[0].size();
    if (!plugin->initialise(ch, have, have)) {
        cerr << "Failed to initialise(" << ch << "," << have << "," << have << ") plugin" << endl;
        return;
    }
    
    float **input = new float *[ch];
    for (int i = 0; i < ch; ++i) {
        input[i] = &samples[i][0];
    }
    Plugin::FeatureSet features = plugin->process(input, RealTime::zeroTime);
    if (!features[0].empty()) {
        fingerprint = QString::fromStdString(features[0][0].label);
    }
    if (!features[1].empty()) {
        puid = QString::fromStdString(features[1][0].label);
    }
    features = plugin->getRemainingFeatures();
    if (fingerprint == "" && !features[0].empty()) {
        fingerprint = QString::fromStdString(features[0][0].label);
    }
    if (puid == "" && !features[1].empty()) {
        puid = QString::fromStdString(features[1][0].label);
    }
    std::cerr << "fingerprint = " << fingerprint << std::endl;
    std::cerr << "puid = " << puid << std::endl;
}

void
guess(QString track)
{
    cout << "Guessing composer for: " << track << endl;

    cerr << "Creating TrackFile object...";
    FileSource fs(track);
    TrackFile *tf = new TrackFile(fs);
    cerr << "done" << endl;
    cerr << "hash = " << tf->hash() << endl;

    QString fingerprint, puid, maker, title;
    AudioFileReader::TagMap tags;
    //!!! bad api!:
    getTrackData(fs, fingerprint, puid, title, maker, tags);

    if (maker != "") {
        cout << "Name found: " << maker << endl;
        if (composerAliases.contains(maker)) {
            QList<Composer *> matching = composerAliases.values(maker);
            foreach (Composer *c, matching) {
                cout << "Matched URI: " << c->uri() << endl;
            }
        } else {
            ComposerFullTextMatcher matcher(allComposers);
            GuessList gl(matcher.match(maker, 5, 0.5));
            if (!gl.empty()) {
                foreach (Guess guess, gl) {
                    cout << "Possibly matched URI (score = " << guess.confidence() << "): " << guess.entity()->uri() << endl;
                }
            }
        }
    }

    if (tags.find("MUSICBRAINZ_ARTISTID") != tags.end()) {
        QString id = tags["MUSICBRAINZ_ARTISTID"];
        Uri mbzUri = Uri("http://dbtune.org/musicbrainz/resource/artist/" + id);
        cout << "MBZ id found: " << id << endl;
        if (composerUris.contains(mbzUri)) {
            cout << "Matched URI: " << composerUris[mbzUri]->uri() << endl;
        }
    }

    
}

int
main(int argc, char **argv)
{
    if (argc < 3) usage(argv[0]);
    QString inRDFFileName = argv[1];
    QString command = argv[2];
    QStringList args;
    for (int i = 3; i < argc; ++i) {
        args.push_back(argv[i]);
    }

    BasicStore *store = new BasicStore();
    store->setBaseUri(Uri("http://dbtune.org/classical/resource/"));
    ObjectLoader *loader = new ObjectLoader(store);

    TypeMapping tm;

    TypeRegistrar::registerTypes();
    TypeRegistrar::addMappings(store, &tm);

    loader->setTypeMapping(tm);

    if (!load(store, inRDFFileName)) {
	cerr << "Failed to load data source" << endl;
	return 1;
    }

    cerr << "Imported RDF data, mapping to objects...";
    QObjectList objects = loader->loadAll();
    cerr << " done" << endl;

    delete loader;
    
    foreach (QObject *o, objects) {
        Composer *c = qobject_cast<Composer *>(o);
        if (c) {
            allComposers.push_back(c);
            composerAliases.insert(c->name(), c);
            foreach (QString alias, c->aliases()) {
                composerAliases.insert(alias, c);
            }
            composerUris.insert(c->property("uri").value<Uri>(), c);
            foreach (Uri otherUri, c->otherURIs()) {
                composerUris.insert(otherUri, c);
            }
        }
    }
    
    QList<Work *> works;
    foreach (QObject *o, objects) {
        Work *w = qobject_cast<Work *>(o);
        if (w) works.push_back(w);
    }

    foreach (Work *w, works) {
        Composition *c = w->composition();
        if (c) {
            Composer *cp = c->composer();
            if (cp) worksMap[cp].insert(w);
        }
    }

    if (command == "guess") {
        if (args.empty()) usage(argv[0]);
        foreach (QString track, args) {
            guess(track);
        }
    } 
        
    delete store;
    TempDirectory::getInstance()->cleanup();
}