added multipolygon to bitmask conversion

Author: czaloom

api/valor_api/backend/core/annotation.py

@@ -1,5 +1,12 @@
-from geoalchemy2.functions import ST_AsGeoJSON
-from sqlalchemy import ScalarSelect, and_, delete, insert, select
+from geoalchemy2.functions import (
+    ST_AddBand,
+    ST_AsGeoJSON,
+    ST_AsRaster,
+    ST_GeomFromText,
+    ST_MakeEmptyRaster,
+    ST_MapAlgebra,
+)
+from sqlalchemy import CTE, ScalarSelect, and_, delete, func, insert, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
@@ -10,6 +17,40 @@
 from valor_api.enums import ModelStatus, TableStatus, TaskType
 
 
+def _create_raster_from_multipolygon(raster: schemas.Raster) -> CTE:
+    if not raster.geometry:
+        raise ValueError("Raster must contain a geometry.")
+    empty_raster = ST_AddBand(
+        ST_MakeEmptyRaster(
+            raster.width,  # width
+            raster.height,  # height
+            0,  # upperleftx
+            0,  # upperlefty
+            1,  # scalex
+            1,  # scaley
+            0,  # skewx
+            0,  # skewy
+            0,  # srid
+        ),
+        "1BB",  # pixeltype
+    )
+    return select(
+        ST_MapAlgebra(
+            empty_raster,
+            ST_AsRaster(
+                ST_GeomFromText(raster.geometry.to_wkt()),
+                empty_raster,
+                "1BB",
+                1,
+                0,
+            ),
+            "[rast2]",
+            "1BB",
+            "UNION",
+        ).label("raster")
+    ).cte()
+
+
 def _format_box(box: schemas.Box | None) -> str | None:
     return box.to_wkt() if box else None
 
@@ -25,10 +66,17 @@ def _format_raster(
 
 
 def _format_bitmask(
+    db: Session,
     raster: schemas.Raster | None,
 ) -> str | None:
-    if not raster or raster.geometry:
+    """
+    Converts a Raster schema into a bitmask.
+    """
+    if raster is None:
         return None
+    elif raster and raster.geometry:
+        r = _create_raster_from_multipolygon(raster)
+        return db.scalar(func.raster_to_bitstring(r.c.raster))
     return "".join("1" if b else "0" for b in raster.array.flatten())
 
 
@@ -112,7 +160,7 @@ def create_annotations(
             "box": _format_box(annotation.bounding_box),
             "polygon": _format_polygon(annotation.polygon),
             "raster": _format_raster(annotation.raster),
-            "bitmask": _format_bitmask(annotation.raster),
+            "bitmask": _format_bitmask(db, annotation.raster),
             "embedding_id": _create_embedding(
                 db=db, value=annotation.embedding
             ),

api/valor_api/backend/metrics/segmentation.py

@@ -143,6 +143,7 @@ def _compute_iou(
         )
 
         if iou is None:
+            print(label)
             continue
 
         metrics.extend(

integration_tests/client/metrics/test_segmentation.py

@@ -34,6 +34,10 @@ def test_evaluate_segmentation(
         model.add_prediction(dataset, pred)
     model.finalize_inferences(dataset)
 
+    import time
+
+    time.sleep(20)
+
     eval_job = model.evaluate_segmentation(dataset)
     assert eval_job.missing_pred_labels == [
         {"key": "k3", "value": "v3", "score": None}

migrations/sql/00000015_bitstring_raster.up.sql

@@ -1 +1,35 @@
-ALTER TABLE annotation ADD COLUMN bitmask bit varying;
+ALTER TABLE annotation ADD COLUMN bitmask bit varying;
+
+CREATE OR REPLACE FUNCTION raster_to_bitstring(input_raster raster)
+RETURNS text AS
+$$
+DECLARE
+    width INTEGER;
+    height INTEGER;
+    bitstring TEXT := '';
+    val INTEGER;
+    pixelval INTEGER;
+BEGIN
+    -- Get the width and height of the raster
+    SELECT ST_Width(input_raster), ST_Height(input_raster) INTO width, height;
+
+    -- Iterate through each row and column of the raster
+    FOR i IN 1..height LOOP
+        FOR j IN 1..width LOOP
+            -- Get the pixel value at position (i, j)
+            pixelval := ST_Value(input_raster, i, j);
+
+            -- 1BB raster has values of 0 or 1, so just append it to the bitstring
+            IF pixelval IS NOT NULL THEN
+                val := pixelval::integer;
+                bitstring := bitstring || val::text;
+            ELSE
+                bitstring := bitstring || '0'; -- or handle NULL as needed
+            END IF;
+        END LOOP;
+    END LOOP;
+
+    RETURN bitstring;
+END;
+$$
+LANGUAGE plpgsql;