Transforms for building AFLW dataset

Jordan Cheney
1 parent 0b967872
Showing 3 changed files with 115 additions and 3 deletions
openbr/plugins/imgproc/rndaffine.cpp
openbr/plugins/io/write.cpp
openbr/plugins/metadata/registerpointsasaffine.cpp
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include <openbr/plugins/openbr_internal.h>
+#include <openbr/core/opencvutils.h>
+
+using namespace cv;
+
+namespace br
+{
+
+class RndAffineTransform : public UntrainableMetaTransform
+{
+    Q_OBJECT
+    Q_PROPERTY(int numAffines READ get_numAffines WRITE set_numAffines RESET reset_numAffines STORED false)
+    Q_PROPERTY(int winWidth READ get_winWidth WRITE set_winWidth RESET reset_winWidth STORED false)
+    Q_PROPERTY(int winHeight READ get_winHeight WRITE set_winHeight RESET reset_winHeight STORED false)
+    Q_PROPERTY(float scaleFactor READ get_scaleFactor WRITE set_scaleFactor RESET reset_scaleFactor STORED false)
+    Q_PROPERTY(int maxAngle READ get_maxAngle WRITE set_maxAngle RESET reset_maxAngle STORED false)
+    BR_PROPERTY(int, numAffines, 0)
+    BR_PROPERTY(int, winWidth, 24)
+    BR_PROPERTY(int, winHeight, 24)
+    BR_PROPERTY(float, scaleFactor, 1.2)
+    BR_PROPERTY(int, maxAngle, 15)
+
+    void project(const Template &src, Template &dst) const
+    {
+        TemplateList temp;
+        project(TemplateList() << src, temp);
+        if (!temp.isEmpty()) dst = temp.first();
+    }
+
+    void project(const TemplateList &src, TemplateList &dst) const
+    {
+        foreach (const Template &t, src) {
+            QPointF affine_0 = t.file.get<QPointF>("Affine_0");
+            QPointF affine_1 = t.file.get<QPointF>("Affine_1");
+
+            // Append the original points
+            Template u = t;
+            u.file.setPoints(QList<QPointF>() << affine_0 << affine_1);
+            u.file.set("Affine_0", affine_0);
+            u.file.set("Affine_1", affine_1);
+            dst.append(u);
+
+            const double IPD = sqrt(pow(affine_0.x() - affine_1.x(), 2) + pow(affine_0.y() - affine_1.y(), 2));
+            for (int i = 0; i < numAffines; i++) {
+                int angle = (rand() % (2*maxAngle)) - maxAngle;
+
+                int min = (int)(sqrt(1 / scaleFactor) * IPD);
+                int max = (int)(sqrt(scaleFactor) * IPD);
+                int dx = (rand() % (max - min)) + min;
+                int dy = (dx * sin(angle * M_PI / 180))/2;
+
+                QPointF shiftedAffine_0 = QPointF(affine_1.x() - dx, affine_1.y() + dy);
+
+                Template u = t;
+                u.file.setPoints(QList<QPointF>() << shiftedAffine_0 << affine_1);
+                u.file.set("Affine_0", shiftedAffine_0);
+                u.file.set("Affine_1", affine_1);
+                dst.append(u);
+            }
+        }
+    }
+};
+
+BR_REGISTER(Transform, RndAffineTransform)
+
+} // namespace br
+
+#include "imgproc/rndaffine.moc"
@@ -31,10 +31,10 @@ class WriteTransform : public TimeVaryingTransform
 {
     Q_OBJECT
     Q_PROPERTY(QString outputDirectory READ get_outputDirectory WRITE set_outputDirectory RESET reset_outputDirectory STORED false)
-    Q_PROPERTY(QString imageName READ get_imageName WRITE set_imageName RESET reset_imageName STORED false)
+    Q_PROPERTY(QString underscore READ get_underscore WRITE set_underscore RESET reset_underscore STORED false)
     Q_PROPERTY(QString imgExtension READ get_imgExtension WRITE set_imgExtension RESET reset_imgExtension STORED false)
     BR_PROPERTY(QString, outputDirectory, "Temp")
-    BR_PROPERTY(QString, imageName, "image")
+    BR_PROPERTY(QString, underscore, "")
     BR_PROPERTY(QString, imgExtension, "jpg")
  
     int cnt;
@@ -48,7 +48,8 @@ class WriteTransform : public TimeVaryingTransform
     void projectUpdate(const Template &src, Template &dst)
     {
         dst = src;
-        OpenCVUtils::saveImage(dst.m(), QString("%1/%2_%3.%4").arg(outputDirectory).arg(imageName).arg(cnt++, 5, 10, QChar('0')).arg(imgExtension));
+        QString path = QString("%1/%2%3.%4").arg(outputDirectory).arg(dst.file.baseName()).arg(underscore.isEmpty() ? "" : "_" + underscore).arg(imgExtension);
+        OpenCVUtils::saveImage(dst.m(), path);
     }
  
 };
+#include <openbr/plugins/openbr_internal.h>
+
+namespace br
+{
+
+class RegisterPointsAsAffine : public UntrainableMetadataTransform
+{
+    Q_OBJECT
+    Q_PROPERTY(QList<int> pointsIdxs READ get_pointIdxs WRITE set_pointIdxs RESET reset_pointIdxs STORED false)
+    BR_PROPERTY(QList<int>, pointIdxs, QList<int>())
+
+    void projectMetadata(const File &src, File &dst) const
+    {
+        const int chin = 20;
+
+        if (pointIdxs.size() != 2 && pointIdxs.size() != 3)
+            qFatal("Need 2 or 3 points for affine transform");
+
+        dst = src;
+
+        QList<QPointF> points = src.points();
+
+        if (points[pointIdxs[0]] == QPointF(-1, -1) ||
+            points[pointIdxs[1]] == QPointF(-1, -1))
+            dst.fte = true;
+
+        if (points[chin] == QPointF(-1, -1))
+            dst.fte = true;
+
+        dst.set("Affine_0", points[pointIdxs[0]]);
+        dst.set("Affine_1", points[pointIdxs[1]]);
+        if (pointIdxs.size() == 3)
+            dst.set("Affine_2", points[pointIdxs[2]]);
+    }
+};
+
+BR_REGISTER(Transform, RegisterPointsAsAffine)
+
+} // namespace br
+
+#include "metadata/registerpointsasaffine.moc"