Added some checks for mask related transforms

Scott Klum
1 parent fd9094e2
Showing 2 changed files with 50 additions and 39 deletions
openbr/plugins/imgproc/cropfrommask.cpp
openbr/plugins/imgproc/samplefrommask.cpp
@@ -24,43 +24,48 @@ private:
         Mat mask = dst.file.get<Mat>("Mask");
-        int w = mask.rows;
-        int h = mask.cols;
-        int left = w;
-        int right = 0;
-        int top = h;
-        int bottom = 0;
-        for (int i = 0 ; i < w; i++) {
-            for (int j = 0 ; j < h; j++) {
-                if (mask.at<unsigned char>(i,j)) {
-                    if (i < left)
-                        left = i;
-                    if (i > right)
-                        right = i;
-                    if (j < top)
-                        top = j;
-                    if (j > bottom)
-                        bottom = j;
+        const int count = countNonZero(mask);
+
+        if (count > 0) {
+            int w = mask.rows;
+            int h = mask.cols;
+            int left = w;
+            int right = 0;
+            int top = h;
+            int bottom = 0;
+            for (int i = 0 ; i < w; i++) {
+                for (int j = 0 ; j < h; j++) {
+                    if (mask.at<unsigned char>(i,j)) {
+                        if (i < left)
+                            left = i;
+                        if (i > right)
+                            right = i;
+                        if (j < top)
+                            top = j;
+                        if (j > bottom)
+                            bottom = j;
+                    }
                 }
             }
-        }
-        if (fixedAspectRatio) {
-            h = bottom - top + 1;
-            w = right - left + 1;
-            if (h > w) {
-                int h2 = (h - w) / 2;
-                right += h2;
-                left -= h2;
-            } else {
-                int w2 = (w - h) / 2;
-                bottom += w2;
-                top -= w2;
+            if (fixedAspectRatio) {
+                h = bottom - top + 1;
+                w = right - left + 1;
+                if (h > w) {
+                    int h2 = (h - w) / 2;
+                    right += h2;
+                    left -= h2;
+                } else {
+                    int w2 = (w - h) / 2;
+                    bottom += w2;
+                    top -= w2;
+                }
             }
-        }
-
-        dst.m() = Mat(src.m(), Rect(top, left, bottom - top + 1, right - left + 1));
+            dst.m() = Mat(src.m(), Rect(top, left, bottom - top + 1, right - left + 1));
+        } else {
+            dst.file.fte = true;
+        }
     }
 };
@@ -18,16 +18,22 @@ class SampleFromMaskTransform : public UntrainableTransform
     {
         Mat mask = src.file.get<Mat>("Mask");
         const int count = countNonZero(mask);
-        dst.m() = Mat(1,count,src.m().type());
-        Mat masked;
-        src.m().copyTo(masked, mask);
+        if (count > 0) {
+            dst.m() = Mat(1,count,src.m().type());
-        Mat indices;
-        findNonZero(masked,indices);
+            Mat masked;
+            src.m().copyTo(masked, mask);
-        for (size_t j=0; j<indices.total(); j++)
-            dst.m().at<uchar>(0,j) = masked.at<uchar>(indices.at<Point>(j).y,indices.at<Point>(j).x);
+            Mat indices;
+            findNonZero(masked,indices);
+
+            for (size_t j=0; j<indices.total(); j++)
+                dst.m().at<uchar>(0,j) = masked.at<uchar>(indices.at<Point>(j).y,indices.at<Point>(j).x);
+        } else {
+            dst.file.fte = true;
+            qWarning("No mask content for %s.",qPrintable(src.file.baseName()));
+        }
     }
 };