Repository: Picterra/picterra-python

Issue:
[beta] Add detector annotations
Add a beta helper function to add annotations to a detector

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 55c0d9f..2944e16 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
 .DS_Store
 .cache
 .vscode
+.venv
 *.egg-info
 .eggs
 __pycache__
diff --git a/examples/data/outlines.geojson b/examples/data/outlines.geojson
new file mode 100644
index 0000000..92ddad7
--- /dev/null
+++ b/examples/data/outlines.geojson
@@ -0,0 +1,36 @@
+{
+  "type": "FeatureCollection",
+  "features": [
+    {
+      "type": "Feature",
+      "properties": {},
+      "geometry": {
+        "type": "Polygon",
+        "coordinates": [
+          [
+            [
+              8.714802861213684,
+              47.74745244535212
+            ],
+            [
+              8.714770674705505,
+              47.74739833982634
+            ],
+            [
+              8.714942336082457,
+              47.747373090561716
+            ],
+            [
+              8.714987933635712,
+              47.74745424886868
+            ],
+            [
+              8.714802861213684,
+              47.74745244535212
+            ]
+          ]
+        ]
+      }
+    }
+  ]
+}
\ No newline at end of file
diff --git a/examples/set_annotations.py b/examples/set_annotations.py
new file mode 100644
index 0000000..7352796
--- /dev/null
+++ b/examples/set_annotations.py
@@ -0,0 +1,12 @@
+import json
+from picterra import APIClient
+
+client = APIClient(api_key='1234')
+# Replace this with the id of one of your detectors
+detector_id = '74576c1f-05d8-48cf-b515-9dba48edb6da'
+# Replace this with the id of one of the training rasters for the detector above
+raster_id = '8331b2a9-c7b1-439e-95e5-647e91abfb8f'
+
+with open('data/outlines.geojson') as f:
+    outlines = json.load(f)
+client.set_annotations(detector_id, raster_id, 'outline', outlines)
diff --git a/src/picterra/client.py b/src/picterra/client.py
index 46a582f..288d55e 100644
--- a/src/picterra/client.py
+++ b/src/picterra/client.py
@@ -47,6 +47,25 @@ class APIClient():
     def _api_url(self, path):
         return urljoin(self.base_url, path)
 
+    def _wait_until_operation_completes(self, operation):
+        operation_id = operation['operation_id']
+        poll_interval = operation['poll_interval']
+        time.sleep(poll_interval * 0.1)
+        while True:
+            logger.info('polling operation id %s' % operation_id)
+            resp = self.sess.get(
+                self._api_url('operations/%s/' % operation_id),
+            )
+            if not resp.ok:
+                raise APIError(resp.text)
+            status = resp.json()['status']
+            logger.info('status=%s' % status)
+            if status == 'success':
+                break
+            if status == 'failed':
+                raise APIError('Operation %s failed' % operation_id)
+            time.sleep(poll_interval)
+
     def upload_raster(self, filename, name):
         """
         Upload a raster to picterra.
@@ -77,25 +96,7 @@ class APIClient():
         resp = self.sess.post(self._api_url('rasters/%s/commit/' % raster_id))
         if not resp.ok:
             raise APIError(resp.text)
-        poll_interval = resp.json()['poll_interval']
-
-        # Wait for raster to be processed
-        def _is_ready():
-            logger.info('checking upload status')
-            resp = self.sess.get(
-                self._api_url('rasters/%s/' % raster_id)
-            )
-            if not resp.ok:
-                logger.warning('failed to get raster status - retrying')
-                return False
-
-            if resp.json()['status'] == 'ready':
-                return True
-            elif resp.json()['status'] == 'failed':
-                raise APIError(resp.text)
-            else:
-                return False
-        _poll_at_interval(_is_ready, poll_interval)
+        self._wait_until_operation_completes(resp.json())
         return raster_id
 
     def list_rasters(self):
@@ -155,25 +156,7 @@ class APIClient():
         )
         if not resp.ok:
             raise APIError(resp.text)
-        poll_interval = resp.json()['poll_interval']
-
-        # Wait for detection area to be associated with raster
-        def _is_ready():
-            logger.info('checking upload status')
-            resp = self.sess.get(
-                self._api_url('rasters/%s/detection_areas/upload/%s/' % (raster_id, upload_id))
-            )
-            if not resp.ok:
-                logger.warning('failed to get detection area status - retrying')
-                return False
-
-            if resp.json()['status'] == 'ready':
-                return True
-            elif resp.json()['status'] == 'failed':
-                raise APIError(resp.text)
-            else:
-                return False
-        _poll_at_interval(_is_ready, poll_interval)
+        self._wait_until_operation_completes(resp.json())
 
     def add_raster_to_detector(self, raster_id: str, detector_id: str) -> str:
         """
@@ -281,3 +264,46 @@ class APIClient():
                 for chunk in r.iter_content(chunk_size=8192):
                     if chunk:  # filter out keep-alive new chunks
                         f.write(chunk)
+
+    def set_annotations(self, detector_id, raster_id, annotation_type, annotations):
+        """
+        Replaces the annotations of type 'annotation_type' with 'annotations', for the
+        given raster-detector pair.
+
+        Args:
+            detector_id (str): The id of the detector
+            raster_id (str): The id of the raster
+            annotation_type (str): One of (outlines, training_area, testing_area, validation_area)
+            annotations (dict): GeoJSON representation of the features to upload
+        """
+        # Get an upload url
+        create_upload_resp = self.sess.post(
+            self._api_url(
+                'detectors/%s/training_rasters/%s/%s/upload/bulk/'
+                % (detector_id, raster_id, annotation_type)
+            )
+        )
+        if not create_upload_resp.ok:
+            raise APIError(create_upload_resp.text)
+
+        upload = create_upload_resp.json()
+        upload_url = upload['upload_url']
+        upload_id = upload['upload_id']
+
+        # Upload data
+        upload_resp = requests.put(upload_url, json=annotations)
+        if not upload_resp.ok:
+            raise APIError(upload_resp.text)
+
+        # Commit upload
+        commit_upload_resp = self.sess.post(
+            self._api_url(
+                'detectors/%s/training_rasters/%s/%s/upload/bulk/%s/commit/'
+                % (detector_id, raster_id, annotation_type, upload_id)
+            )
+        )
+        if not commit_upload_resp.ok:
+            raise APIError(commit_upload_resp.text)
+
+        # Poll for operation completion
+        self._wait_until_operation_completes(commit_upload_resp.json())


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 135bbea..d46a8ae 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -37,6 +37,52 @@ def add_mock_detector_creation_response():
     responses.add(responses.POST, api_url('detectors/'), json={'id': 'foobar'}, status=201)
 
 
+def add_mock_operations_responses(status):
+    operation_id = 21
+    data = {
+        'type': 'mock_operation_type',
+        'status': status
+    }
+    responses.add(
+        responses.GET, 
+        api_url('operations/%s/' % operation_id),
+        json=data, status=200
+    )
+
+
+def add_mock_annotations_responses(detector_id, raster_id):
+    upload_id = 32
+    operation_id = 21
+
+    responses.add(responses.PUT, 'http://storage.example.com', status=200)
+    
+    for annotation_type in ('outline', 'training_area', 'testing_area', 'validation_area'):
+        responses.add(
+            responses.POST,
+            api_url(
+                'detectors/%s/training_rasters/%s/%s/upload/bulk/'
+                % (detector_id, raster_id, annotation_type)
+            ),
+            json={
+                'upload_url': 'http://storage.example.com',
+                'upload_id': upload_id
+            },
+            status=201
+        )
+        responses.add(
+            responses.POST,
+            api_url(
+                'detectors/%s/training_rasters/%s/%s/upload/bulk/%s/commit/'
+                % (detector_id, raster_id, annotation_type, upload_id)
+            ),
+            json={
+                'operation_id': operation_id,
+                'poll_interval': TEST_POLL_INTERVAL
+            },
+            status=201
+        )
+
+
 def add_mock_raster_upload_responses():
     raster_id = 42
     # Upload initiation
@@ -54,6 +100,7 @@ def add_mock_raster_upload_responses():
     # Commit
     data = {
         'poll_interval': TEST_POLL_INTERVAL,
+        'operation_id': 21
     }
     responses.add(
         responses.POST,
@@ -85,6 +132,7 @@ def add_mock_raster_upload_responses():
 
 def add_mock_detection_areas_upload_responses(raster_id):
     upload_id = 42
+    
     # Upload initiation
     data = {
         'upload_url': 'http://storage.example.com',
@@ -100,6 +148,7 @@ def add_mock_detection_areas_upload_responses(raster_id):
     # Commit
     data = {
         'poll_interval': TEST_POLL_INTERVAL,
+        'operation_id': 21
     }
     responses.add(
         responses.POST,
@@ -169,6 +218,7 @@ def add_mock_download_result_response(result_id):
 def test_upload_raster():
     client = _client()
     add_mock_raster_upload_responses()
+    add_mock_operations_responses('success')
     # This just tests that this doesn't raise
     with tempfile.NamedTemporaryFile() as f:
         client.upload_raster(f.name, name='test 1')
@@ -194,6 +244,7 @@ def test_detector_creation():
 @responses.activate
 def test_set_raster_detection_areas_from_file():
     add_mock_detection_areas_upload_responses(1)
+    add_mock_operations_responses('success')
 
     client = _client()
     # This just tests that this doesn't raise
@@ -217,3 +268,14 @@ def test_download_result_to_file():
     with tempfile.NamedTemporaryFile() as f:
         client.download_result_to_file(101, f.name)
         assert open(f.name).read() == expected_content
+
+
+@responses.activate
+def test_upload_annotations():
+    add_mock_annotations_responses(1, 2)
+    add_mock_operations_responses('pending')
+    add_mock_operations_responses('running')
+    add_mock_operations_responses('success')
+
+    client = _client()
+    client.set_annotations(1, 2, 'outline', {})
\ No newline at end of file