diff transform/CSVFeatureWriter.cpp @ 1001:51bf067de517

Add fill-ends option to CSV writer (and test it)

--- a/transform/CSVFeatureWriter.cpp	Tue Oct 14 17:30:37 2014 +0100
+++ b/transform/CSVFeatureWriter.cpp	Wed Oct 15 10:18:13 2014 +0100
@@ -34,7 +34,8 @@
                       "csv"),
     m_separator(","),
     m_sampleTiming(false),
-    m_endTimes(false)
+    m_endTimes(false),
+    m_forceEnd(false)
 {
 }
 
@@ -69,6 +70,11 @@
     p.hasArg = false;
     pl.push_back(p);
 
+    p.name = "fill-ends";
+    p.description = "Include durations (or end times) even for features without duration, by using the gap to the next feature instead.";
+    p.hasArg = false;
+    pl.push_back(p);
+
     return pl;
 }
 
@@ -87,6 +93,8 @@
             m_sampleTiming = true;
         } else if (i->first == "end-times") {
             m_endTimes = true;
+        } else if (i->first == "fill-ends") {
+            m_forceEnd = true;
         }
     }
 }
@@ -98,79 +106,155 @@
                         const Plugin::FeatureList& features,
                         std::string summaryType)
 {
+    TransformId transformId = transform.getIdentifier();
+
     // Select appropriate output file for our track/transform
     // combination
 
-    QTextStream *sptr = getOutputStream(trackId, transform.getIdentifier());
+    QTextStream *sptr = getOutputStream(trackId, transformId);
     if (!sptr) {
-        throw FailedToOpenOutputStream(trackId, transform.getIdentifier());
+        throw FailedToOpenOutputStream(trackId, transformId);
     }
 
     QTextStream &stream = *sptr;
 
-    for (unsigned int i = 0; i < features.size(); ++i) {
+    int n = features.size();
 
-        if (m_stdout || m_singleFileName != "") {
-            if (trackId != m_prevPrintedTrackId) {
-                stream << "\"" << trackId << "\"" << m_separator;
-                m_prevPrintedTrackId = trackId;
+    if (n == 0) return;
+
+    TrackTransformPair tt(trackId, transformId);
+
+    if (m_rates.find(transformId) == m_rates.end()) {
+        m_rates[transformId] = transform.getSampleRate();
+    }
+
+    if (m_pending.find(tt) != m_pending.end()) {
+        writeFeature(tt,
+                     stream,
+                     m_pending[tt],
+                     &features[0],
+                     m_pendingSummaryTypes[tt]);
+        m_pending.erase(tt);
+        m_pendingSummaryTypes.erase(tt);
+    }
+
+    if (m_forceEnd) {
+        // can't write final feature until we know its end time
+        --n;
+        m_pending[tt] = features[n];
+        m_pendingSummaryTypes[tt] = summaryType;
+    }
+
+    for (int i = 0; i < n; ++i) {
+        writeFeature(tt, 
+                     stream,
+                     features[i], 
+                     m_forceEnd ? &features[i+1] : 0,
+                     summaryType);
+    }
+}
+
+void
+CSVFeatureWriter::finish()
+{
+    for (PendingFeatures::const_iterator i = m_pending.begin();
+         i != m_pending.end(); ++i) {
+        TrackTransformPair tt = i->first;
+        Plugin::Feature f = i->second;
+        QTextStream *sptr = getOutputStream(tt.first, tt.second);
+        if (!sptr) {
+            throw FailedToOpenOutputStream(tt.first, tt.second);
+        }
+        QTextStream &stream = *sptr;
+        // final feature has its own time as end time (we can't
+        // reliably determine the end of audio file, and because of
+        // the nature of block processing, the feature could even
+        // start beyond that anyway)
+        writeFeature(tt, stream, f, &f, m_pendingSummaryTypes[tt]);
+    }
+
+    m_pending.clear();
+}
+
+void
+CSVFeatureWriter::writeFeature(TrackTransformPair tt,
+                               QTextStream &stream,
+                               const Plugin::Feature &f,
+                               const Plugin::Feature *optionalNextFeature,
+                               std::string summaryType)
+{
+    QString trackId = tt.first;
+    TransformId transformId = tt.second;
+
+    if (m_stdout || m_singleFileName != "") {
+        if (trackId != m_prevPrintedTrackId) {
+            stream << "\"" << trackId << "\"" << m_separator;
+            m_prevPrintedTrackId = trackId;
+        } else {
+            stream << m_separator;
+        }
+    }
+
+    Vamp::RealTime duration;
+    bool haveDuration = true;
+    
+    if (f.hasDuration) {
+        duration = f.duration;
+    } else if (optionalNextFeature) {
+        duration = optionalNextFeature->timestamp - f.timestamp;
+    } else {
+        haveDuration = false;
+    }
+
+    if (m_sampleTiming) {
+
+        float rate = m_rates[transformId];
+
+        stream << Vamp::RealTime::realTime2Frame(f.timestamp, rate);
+
+        if (haveDuration) {
+            stream << m_separator;
+            if (m_endTimes) {
+                stream << Vamp::RealTime::realTime2Frame
+                    (f.timestamp + duration, rate);
             } else {
-                stream << m_separator;
+                stream << Vamp::RealTime::realTime2Frame(duration, rate);
             }
         }
 
-        if (m_sampleTiming) {
+    } else {
 
-            stream << Vamp::RealTime::realTime2Frame
-                (features[i].timestamp, transform.getSampleRate());
+        QString timestamp = f.timestamp.toString().c_str();
+        timestamp.replace(QRegExp("^ +"), "");
+        stream << timestamp;
 
-            if (features[i].hasDuration) {
-                stream << m_separator;
-                if (m_endTimes) {
-                    stream << Vamp::RealTime::realTime2Frame
-                        (features[i].timestamp + features[i].duration,
-                         transform.getSampleRate());
-                } else {
-                    stream << Vamp::RealTime::realTime2Frame
-                        (features[i].duration, transform.getSampleRate());
-                }
+        if (haveDuration) {
+            if (m_endTimes) {
+                QString endtime =
+                    (f.timestamp + duration).toString().c_str();
+                endtime.replace(QRegExp("^ +"), "");
+                stream << m_separator << endtime;
+            } else {
+                QString d = duration.toString().c_str();
+                d.replace(QRegExp("^ +"), "");
+                stream << m_separator << d;
             }
+        }            
+    }
 
-        } else {
-
-            QString timestamp = features[i].timestamp.toString().c_str();
-            timestamp.replace(QRegExp("^ +"), "");
-            stream << timestamp;
-
-            if (features[i].hasDuration) {
-                if (m_endTimes) {
-                    QString endtime =
-                        (features[i].timestamp + features[i].duration)
-                        .toString().c_str();
-                    endtime.replace(QRegExp("^ +"), "");
-                    stream << m_separator << endtime;
-                } else {
-                    QString duration = features[i].duration.toString().c_str();
-                    duration.replace(QRegExp("^ +"), "");
-                    stream << m_separator << duration;
-                }
-            }            
-        }
-
-        if (summaryType != "") {
-            stream << m_separator << summaryType.c_str();
-        }
-
-        for (unsigned int j = 0; j < features[i].values.size(); ++j) {
-            stream << m_separator << features[i].values[j];
-        }
-
-        if (features[i].label != "") {
-            stream << m_separator << "\"" << features[i].label.c_str() << "\"";
-        }
-
-        stream << "\n";
+    if (summaryType != "") {
+        stream << m_separator << summaryType.c_str();
     }
+    
+    for (unsigned int j = 0; j < f.values.size(); ++j) {
+        stream << m_separator << f.values[j];
+    }
+    
+    if (f.label != "") {
+        stream << m_separator << "\"" << f.label.c_str() << "\"";
+    }
+    
+    stream << "\n";
 }