migrating to bitstring rasters

Author: czaloom

api/tests/functional-tests/backend/core/test_geometry.py

@@ -11,7 +11,6 @@
     _convert_polygon_to_box,
     _convert_raster_to_box,
     _convert_raster_to_polygon,
-    _raster_to_png_b64,
     convert_geometry,
     get_annotation_type,
 )
@@ -206,7 +205,7 @@ def test_convert_from_raster(
             and_(
                 models.Annotation.box.is_(None),
                 models.Annotation.polygon.is_(None),
-                models.Annotation.raster.isnot(None),
+                models.Annotation.bitmask_id.isnot(None),
             )
         )
     )
@@ -227,7 +226,7 @@ def test_convert_from_raster(
 
     assert annotation.box is not None
     assert annotation.polygon is not None
-    assert annotation.raster is not None
+    assert annotation.bitmask_id is not None
 
     converted_box = _load_box(db, annotation.box)
     converted_polygon = _load_polygon(db, annotation.polygon)
@@ -257,7 +256,7 @@ def test_convert_polygon_to_box(
             and_(
                 models.Annotation.box.is_(None),
                 models.Annotation.polygon.isnot(None),
-                models.Annotation.raster.is_(None),
+                models.Annotation.bitmask_id.is_(None),
             )
         )
     )
@@ -275,7 +274,7 @@ def test_convert_polygon_to_box(
 
     assert annotation.box is not None
     assert annotation.polygon is not None
-    assert annotation.raster is None
+    assert annotation.bitmask_id is None
 
     converted_box = _load_box(db, annotation.box)
 
@@ -317,8 +316,8 @@ def test_create_raster_from_polygons(
 
     # verify all rasters are equal
     raster_arrs = [
-        Raster(mask=_raster_to_png_b64(db, r)).to_numpy()
-        for r in db.scalars(select(models.Annotation.raster)).all()
+        np.array([bit == "1" for bit in r.value]).reshape((r.height, r.width))
+        for r in db.query(models.Bitmask).all()
     ]
     assert len(raster_arrs) == 3
 

api/tests/functional-tests/backend/metrics/test_segmentation.py

@@ -199,7 +199,10 @@ def test__count_true_positives(
 
         expected = _help_count_true_positives(
             gt_semantic_segs_create,
-            [pred_semantic_segs_img1_create, pred_semantic_segs_img2_create],
+            [
+                pred_semantic_segs_img1_create,
+                pred_semantic_segs_img2_create,
+            ],
             schemas.Label(key=k, value=v),
         )
 

api/tests/functional-tests/crud/test_create_delete.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import pytest
-from geoalchemy2.functions import ST_AsText, ST_Count, ST_Polygon
+from geoalchemy2.functions import ST_AsText, ST_Polygon
 from PIL import Image
 from sqlalchemy import func, select
 from sqlalchemy.orm import Session
@@ -815,9 +815,12 @@ def test_create_predicted_segmentations_check_area_and_delete_model(
 
     raster_counts = set(
         db.scalars(
-            select(ST_Count(models.Annotation.raster)).where(
-                models.Annotation.model_id.isnot(None)
+            select(func.bit_count(models.Bitmask.value))
+            .join(
+                models.Annotation,
+                models.Annotation.bitmask_id == models.Bitmask.id,
             )
+            .where(models.Annotation.model_id.isnot(None))
         )
     )
 
@@ -879,7 +882,9 @@ def test_segmentation_area_no_hole(
         ],
     )
 
-    segmentation_count = db.scalar(select(ST_Count(models.Annotation.raster)))
+    segmentation_count = db.scalar(
+        select(func.bit_count(models.Bitmask.value))
+    )
 
     assert segmentation_count == math.ceil(45.5)  # area of mask will be an int
 
@@ -920,11 +925,9 @@ def test_segmentation_area_with_hole(
         ],
     )
 
-    segmentation = db.scalar(select(models.Annotation))
-
     # give tolerance of 2 pixels because of poly -> mask conversion
-    assert segmentation
-    assert (db.scalar(ST_Count(segmentation.raster)) - 92) <= 2
+    assert db.scalar(select(models.Annotation))
+    assert (db.scalar(func.bit_count(models.Bitmask.value)) - 92) <= 2
 
 
 def test_segmentation_area_multi_polygon(
@@ -966,13 +969,14 @@ def test_segmentation_area_multi_polygon(
         ],
     )
 
-    segmentation = db.scalar(select(models.Annotation))
-
     # the two shapes don't intersect so area should be sum of the areas
     # give tolerance of 2 pixels because of poly -> mask conversion
-    assert segmentation
+    assert db.scalar(select(models.Annotation))
     assert (
-        abs(db.scalar(ST_Count(segmentation.raster)) - (math.ceil(45.5) + 92))
+        abs(
+            db.scalar(func.bit_count(models.Bitmask.value))
+            - (math.ceil(45.5) + 92)
+        )
         <= 2
     )
 
@@ -1037,7 +1041,15 @@ def test_gt_seg_as_mask_or_polys(
 
     shapes = db.scalars(
         select(
-            ST_AsText(ST_Polygon(models.Annotation.raster)),
+            ST_AsText(
+                ST_Polygon(
+                    func.bitstring_to_raster(
+                        models.Bitmask.value,
+                        models.Bitmask.height,
+                        models.Bitmask.width,
+                    )
+                )
+            ),
         )
     ).all()
     assert len(shapes) == 2

api/valor_api/backend/core/annotation.py

@@ -1,3 +1,4 @@
+import numpy as np
 from geoalchemy2.functions import (
     ST_AddBand,
     ST_AsGeoJSON,
@@ -6,13 +7,12 @@
     ST_MakeEmptyRaster,
     ST_MapAlgebra,
 )
-from sqlalchemy import CTE, ScalarSelect, and_, delete, func, insert, select
+from sqlalchemy import CTE, and_, delete, func, insert, select
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
 from valor_api import schemas
 from valor_api.backend import models
-from valor_api.backend.core.geometry import _raster_to_png_b64
 from valor_api.backend.query import generate_query
 from valor_api.enums import ModelStatus, TableStatus, TaskType
 
@@ -59,25 +59,39 @@ def _format_polygon(polygon: schemas.Polygon | None) -> str | None:
     return polygon.to_wkt() if polygon else None
 
 
-def _format_raster(
-    raster: schemas.Raster | None,
-) -> ScalarSelect | bytes | None:
-    return raster.to_psql() if raster else None
+# def _format_raster(
+#     raster: schemas.Raster | None,
+# ) -> ScalarSelect | bytes | None:
+#     return raster.to_psql() if raster else None
 
 
-def _format_bitmask(
+def _create_bitmask(
     db: Session,
     raster: schemas.Raster | None,
-) -> str | None:
+) -> int | None:
     """
-    Converts a Raster schema into a bitmask.
+    Creates a bitmask from a raster schema.
     """
     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())
+        bitstring = db.scalar(func.raster_to_bitstring(r.c.raster))
+    else:
+        bitstring = "".join("1" if b else "0" for b in raster.array.flatten())
+
+    try:
+        row = models.Bitmask(
+            value=bitstring,
+            height=raster.height,
+            width=raster.width,
+        )
+        db.add(row)
+        db.commit()
+    except IntegrityError as e:
+        db.rollback()
+        raise e
+    return row.id
 
 
 def _create_embedding(
@@ -159,8 +173,7 @@ def create_annotations(
             "meta": annotation.metadata,
             "box": _format_box(annotation.bounding_box),
             "polygon": _format_polygon(annotation.polygon),
-            "raster": _format_raster(annotation.raster),
-            "bitmask": _format_bitmask(db, annotation.raster),
+            "bitmask_id": _create_bitmask(db, annotation.raster),
             "embedding_id": _create_embedding(
                 db=db, value=annotation.embedding
             ),
@@ -222,8 +235,7 @@ def create_skipped_annotations(
             meta=dict(),
             box=None,
             polygon=None,
-            raster=None,
-            bitmask=None,
+            bitmask_id=None,
             embedding_id=None,
             text=None,
             context_list=None,
@@ -305,16 +317,19 @@ def get_annotation(
         )
 
     # raster
-    if annotation.raster is not None:
-        datum = db.scalar(
-            select(models.Datum).where(models.Datum.id == annotation.datum_id)
+    if annotation.bitmask_id is not None:
+        bitmask = (
+            db.query(models.Bitmask)
+            .where(models.Bitmask.id == annotation.bitmask_id)
+            .scalar()
         )
-        if datum is None:
-            raise RuntimeError(
-                "psql unexpectedly returned None instead of a Datum."
-            )
-        raster = schemas.Raster(
-            mask=_raster_to_png_b64(db=db, raster=annotation.raster),
+        if bitmask is None:
+            raise ValueError("Expected bitmask to contain a value.")
+        raster = schemas.Raster.from_numpy(
+            mask=np.array(
+                [int(bit) == 1 for bit in bitmask.value],
+                dtype=bool,
+            ).reshape(bitmask.height, bitmask.width)
         )
 
     # embedding
@@ -381,6 +396,28 @@ def get_annotations(
     ]
 
 
+def _delete_linked_annotations(db: Session):
+    # delete embeddings that are no longer referenced
+    existing_ids = select(models.Annotation.embedding_id).where(
+        models.Annotation.embedding_id.isnot(None)
+    )
+    db.execute(
+        delete(models.Embedding).where(
+            models.Embedding.id.not_in(existing_ids)
+        )
+    )
+    db.commit()
+
+    # delete bitmasks that are no longer referenced
+    existing_ids = select(models.Annotation.bitmask_id).where(
+        models.Annotation.bitmask_id.isnot(None)
+    )
+    db.execute(
+        delete(models.Bitmask).where(models.Bitmask.id.not_in(existing_ids))
+    )
+    db.commit()
+
+
 def delete_dataset_annotations(
     db: Session,
     dataset: models.Dataset,
@@ -421,16 +458,9 @@ def delete_dataset_annotations(
         )
         db.commit()
 
-        # delete embeddings (if they exist)
-        existing_ids = select(models.Annotation.embedding_id).where(
-            models.Annotation.embedding_id.isnot(None)
-        )
-        db.execute(
-            delete(models.Embedding).where(
-                models.Embedding.id.not_in(existing_ids)
-            )
-        )
-        db.commit()
+        # delete linked annotations
+        _delete_linked_annotations(db=db)
+
     except IntegrityError as e:
         db.rollback()
         raise e
@@ -475,16 +505,9 @@ def delete_model_annotations(
         )
         db.commit()
 
-        # delete embeddings (if they exist)
-        existing_ids = select(models.Annotation.embedding_id).where(
-            models.Annotation.embedding_id.isnot(None)
-        )
-        db.execute(
-            delete(models.Embedding).where(
-                models.Embedding.id.not_in(existing_ids)
-            )
-        )
-        db.commit()
+        # delete linked annotations
+        _delete_linked_annotations(db=db)
+
     except IntegrityError as e:
         db.rollback()
         raise e

api/valor_api/backend/core/dataset.py

@@ -375,7 +375,7 @@ def get_n_groundtruth_rasters_in_dataset(db: Session, name: str) -> int:
             and_(
                 models.Dataset.name == name,
                 models.Dataset.status != enums.TableStatus.DELETING,
-                models.Annotation.raster.isnot(None),
+                models.Annotation.bitmask_id.isnot(None),
             )
         )
         .distinct()