From ae5c944e3eefc70b04da5f05df07a7e44ff24e71 Mon Sep 17 00:00:00 2001
From: Sergiy <sergiy.popovich@gmail.com>
Date: Sat, 22 Feb 2025 15:54:25 +0000
Subject: [PATCH] feat: deterministic annotatino ids

---
 zetta_utils/db_annotations/annotation.py | 78 +++++++++++++++++++-----
 1 file changed, 62 insertions(+), 16 deletions(-)

diff --git a/zetta_utils/db_annotations/annotation.py b/zetta_utils/db_annotations/annotation.py
index 9a866abe5..4199264a4 100644
--- a/zetta_utils/db_annotations/annotation.py
+++ b/zetta_utils/db_annotations/annotation.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import copy
+import json
 import time
 import uuid
 from typing import Any, Union, cast, overload
@@ -14,7 +15,10 @@
     LineAnnotation,
     PointAnnotation,
 )
+from typeguard import typechecked
 
+from zetta_utils import builder
+from zetta_utils.geometry.bbox import BBox3D
 from zetta_utils.layer.db_layer import DBRowDataT
 from zetta_utils.layer.db_layer.firestore import build_firestore_layer
 from zetta_utils.parsing.ngl_state import AnnotationKeys
@@ -135,23 +139,42 @@ def read_annotations(
         return {k: AnnotationDBEntry.from_dict(k, cast(dict, v)) for k, v in result_raw.items()}
 
 
-def add_annotation(
+def _make_annotation_row(
     annotation: NgAnnotation,
-    *,
     collection_id: str,
     layer_group_id: str,
-    comment: str | None = None,
-    tags: list[str] | None = None,
-) -> str:
+    comment: str | None,
+    tags: list[str] | None,
+) -> tuple[str, dict]:
     row = annotation.to_json()
     row["collection"] = collection_id
     row["layer_group"] = layer_group_id
     row["comment"] = comment
-    row["created_at"] = time.time()
 
+    row_str = json.dumps(row, sort_keys=True)
+    annotation_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, row_str))
+    row["created_at"] = time.time()
     if tags:
         row["tags"] = list(set(tags))
-    annotation_id = str(uuid.uuid4())
+    print(annotation_id)
+    return annotation_id, row
+
+
+def add_annotation(
+    annotation: NgAnnotation,
+    *,
+    collection_id: str,
+    layer_group_id: str,
+    comment: str | None = None,
+    tags: list[str] | None = None,
+) -> str:
+    annotation_id, row = _make_annotation_row(
+        annotation,
+        collection_id=collection_id,
+        layer_group_id=layer_group_id,
+        comment=comment,
+        tags=tags,
+    )
     col_keys = INDEXED_COLS + NON_INDEXED_COLS
     ANNOTATIONS_DB[(annotation_id, col_keys)] = row
     return annotation_id
@@ -167,16 +190,16 @@ def add_annotations(
 ) -> list[str]:
     rows = []
     annotation_ids = []
-    for ann in annotations:
-        row = ann.to_json()
-        row["collection"] = collection_id
-        row["layer_group"] = layer_group_id
-        row["comment"] = comment
-        row["created_at"] = time.time()
-        if tags:
-            row["tags"] = list(set(tags))
+    for annotation in annotations:
+        annotation_id, row = _make_annotation_row(
+            annotation,
+            collection_id=collection_id,
+            layer_group_id=layer_group_id,
+            comment=comment,
+            tags=tags,
+        )
         rows.append(row)
-        annotation_ids.append(str(uuid.uuid4()))
+        annotation_ids.append(annotation_id)
     col_keys = INDEXED_COLS + NON_INDEXED_COLS
     ANNOTATIONS_DB[(annotation_ids, col_keys)] = rows
     return annotation_ids
@@ -250,3 +273,26 @@ def parse_ng_annotations(annotations_raw: list[dict]) -> list[NgAnnotation]:
         else:
             annotations.append(AxisAlignedBoundingBoxAnnotation(**ann))
     return annotations
+
+
+@typechecked
+@builder.register("add_bbox_annotation")
+def add_bbox_annotation(
+    bbox: BBox3D,
+    *,
+    collection_id: str,
+    layer_group_id: str,
+    comment: str | None = None,
+    tags: list[str] | None = None,
+) -> str:
+    ng_bbox = AxisAlignedBoundingBoxAnnotation(
+        point_a=bbox.start,
+        point_b=bbox.end,
+    )
+    return add_annotation(
+        ng_bbox,
+        collection_id=collection_id,
+        layer_group_id=layer_group_id,
+        comment=comment,
+        tags=tags,
+    )