Mask on LazyMultibandRasters directly

geotrellis · Apr 10, 2019 · 99bfe66 · 99bfe66
1 parent 0b7dd6a
commit 99bfe66
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/jvm/src/main/scala/eval/directive/OpDirectives.scala b/jvm/src/main/scala/eval/directive/OpDirectives.scala
@@ -360,7 +360,7 @@ object OpDirectives {
     }).andThen({ case (lzRaster, geom) =>
       geom.as[MultiPolygon] match {
         case Some(mp) =>
-          Valid(ImageResult(LazyMultibandRaster(List(MaskingNode(lzRaster.bands.values.toList, mp)))))
+          Valid(ImageResult(lzRaster.mask(mp)))
         case None =>
           Invalid(NEL.of(NonEvaluableNode(mask, Some("Masking operation requires its vector argument to be a multipolygon"))))
       }

diff --git a/jvm/src/main/scala/eval/tile/LazyMultibandRaster.scala b/jvm/src/main/scala/eval/tile/LazyMultibandRaster.scala
@@ -87,6 +87,13 @@ case class LazyMultibandRaster(val bands: Map[String, LazyRaster]) {
     val lztiles = bands.mapValues({ lt => LazyRaster.Hillshade(List(lt), gridbounds, zFactor, cs, azimuth, altitude) })
     LazyMultibandRaster(lztiles)
   }
+
+  def mask(
+    maskPoly: MultiPolygon
+  ): LazyMultibandRaster = {
+    val lztiles = bands.mapValues({ lt => MaskingNode(List(lt), maskPoly) })
+    LazyMultibandRaster(lztiles)
+  }
 }
 
 object LazyMultibandRaster {

diff --git a/jvm/src/test/scala/eval/ResultSpec.scala b/jvm/src/test/scala/eval/ResultSpec.scala
@@ -62,15 +62,23 @@ class ResultSpec extends FunSpec with Matchers {
     val mask = Extent(0, 0, 1, 1).as[Polygon].map(MultiPolygon(_)).get
     val maskResult = ImageResult(LazyMultibandRaster(List(MaskingNode(List(rasterOnes), mask))))
 
+
+    val maskResultSB = ImageResult(LazyMultibandRaster(List(MaskingNode(List(rasterOnes), mask))))
+    val maskResultRGB = ImageResult(LazyMultibandRaster(
+                                      List(rasterOnes, rasterOnes, rasterOnes)).mask(mask))
+
     for {
       x <- (0 to 3 toArray)
       y <- (0 to 3 toArray)
     } yield {
-      val fetched = maskResult.res.bands.head._2.get(x, y)
+      val fetchedSB = maskResultSB.res.bands.head._2.get(x, y)
+      val fetchedRGB = maskResultRGB.res.bands.toList map { _._2.get(x, y) }
       if ((x, y) == (0, 3)) {
-        isData(fetched) should be (true)
+        isData(fetchedSB) should be (true)
+        fetchedRGB map { isData(_) } should be (List(true, true, true))
       } else {
-        isData(fetched) should be (false)
+        isData(fetchedSB) should be (false)
+        fetchedRGB map { isData(_) } should be (List(false, false, false))
       }
     }
   }