Merge pull request #41 from scaleapi/metadata-schema

Metadata schema

Author: Nastia39

nucleus/__init__.py

@@ -116,6 +116,9 @@
     NAME_KEY,
     ANNOTATIONS_KEY,
     AUTOTAGS_KEY,
+    ANNOTATION_METADATA_SCHEMA_KEY,
+    ITEM_METADATA_SCHEMA_KEY,
+    FORCE_KEY,
 )
 from .model import Model
 from .errors import (
@@ -150,7 +153,16 @@ def list_models(self) -> List[Model]:
         """
         model_objects = self._make_request({}, "models/", requests.get)
 
-        return [Model(model["id"], model["name"], model["ref_id"], model["metadata"], self) for model in model_objects["models"]]
+        return [
+            Model(
+                model["id"],
+                model["name"],
+                model["ref_id"],
+                model["metadata"],
+                self,
+            )
+            for model in model_objects["models"]
+        ]
 
     def list_datasets(self) -> Dict[str, Union[str, List[str]]]:
         """
@@ -229,15 +241,28 @@ def create_dataset_from_project(
         response = self._make_request(payload, "dataset/create_from_project")
         return Dataset(response[DATASET_ID_KEY], self)
 
-    def create_dataset(self, name: str) -> Dataset:
+    def create_dataset(
+        self,
+        name: str,
+        item_metadata_schema: Optional[Dict] = None,
+        annotation_metadata_schema: Optional[Dict] = None,
+    ) -> Dataset:
         """
-        Creates a new dataset based on payload params:
-        name -- A human-readable name of the dataset.
+        Creates a new dataset:
         Returns a response with internal id and name for a new dataset.
-        :param payload: { "name": str }
+        :param name -- A human-readable name of the dataset.
+        :param item_metadata_schema -- optional dictionary to define item metadata schema
+        :param annotation_metadata_schema -- optional dictionary to define annotation metadata schema
         :return: new Dataset object
         """
-        response = self._make_request({NAME_KEY: name}, "dataset/create")
+        response = self._make_request(
+            {
+                NAME_KEY: name,
+                ANNOTATION_METADATA_SCHEMA_KEY: annotation_metadata_schema,
+                ITEM_METADATA_SCHEMA_KEY: item_metadata_schema,
+            },
+            "dataset/create",
+        )
         return Dataset(response[DATASET_ID_KEY], self)
 
     def delete_dataset(self, dataset_id: str) -> dict:
@@ -325,16 +350,16 @@ def populate_dataset(
         async_responses: List[Any] = []
 
         for batch in tqdm_local_batches:
-            payload = construct_append_payload(batch)
+            payload = construct_append_payload(batch, force)
             responses = self._process_append_requests_local(
-                dataset_id, payload
+                dataset_id, payload, force
             )
             async_responses.extend(responses)
 
         for batch in tqdm_remote_batches:
-            payload = construct_append_payload(batch)
+            payload = construct_append_payload(batch, force)
             responses = self._process_append_requests(
-                dataset_id, payload, batch_size, batch_size
+                dataset_id, payload, force, batch_size, batch_size
             )
             async_responses.extend(responses)
 
@@ -411,7 +436,6 @@ def close_files(request_items):
         # don't forget to close all open files
         for p in request_payloads:
             close_files(p)
-        # [close_files(p) for p in request_payloads]
 
         # response object will be None if an error occurred
         async_responses = [
@@ -428,6 +452,7 @@ def _process_append_requests(
         self,
         dataset_id: str,
         payload: dict,
+        update: bool,
         batch_size: int = 20,
         size: int = 10,
     ):
@@ -446,7 +471,7 @@ def exception_handler(request, exception):
         items = payload[ITEMS_KEY]
         payloads = [
             # batch_size images per request
-            {ITEMS_KEY: items[i : i + batch_size]}
+            {ITEMS_KEY: items[i : i + batch_size], FORCE_KEY: update}
             for i in range(0, len(items), batch_size)
         ]
 
@@ -479,7 +504,7 @@ def annotate_dataset(
             Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
         ],
         update: bool,
-        batch_size: int = 100,
+        batch_size: int = 5000,
     ):
         """
         Uploads ground truth annotations for a given dataset.
@@ -1009,7 +1034,7 @@ def _make_grequest(
 
     def _make_request_raw(
         self, payload: dict, route: str, requests_command=requests.post
-    ) -> dict:
+    ):
         """
         Makes a request to Nucleus endpoint. This method returns the raw
         requests.Response object which is useful for unit testing.
@@ -1046,7 +1071,7 @@ def _make_request(
         """
         response = self._make_request_raw(payload, route, requests_command)
 
-        if response.status_code not in SUCCESS_STATUS_CODES:
+        if getattr(response, "status_code") not in SUCCESS_STATUS_CODES:
             logger.warning(response)
 
         return (

nucleus/constants.py

@@ -33,7 +33,7 @@
 DATASET_MODEL_RUNS_KEY = "model_run_ids"
 DATASET_SLICES_KEY = "slice_ids"
 DATASET_LENGTH_KEY = "length"
-FORCE_KEY = "force"
+FORCE_KEY = "update"
 METADATA_KEY = "metadata"
 NAME_KEY = "name"
 LABEL_KEY = "label"
@@ -51,6 +51,8 @@
 ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 GEOMETRY_KEY = "geometry"
 AUTOTAGS_KEY = "autotags"
+ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
+ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 MASK_URL_KEY = "mask_url"
 INDEX_KEY = "index"
 SEGMENTATIONS_KEY = "segmentations"

nucleus/dataset.py

@@ -75,6 +75,7 @@ def create_model_run(
         reference_id: Optional[str] = None,
         model_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        annotation_metadata_schema: Optional[Dict] = None,
     ):
         """
         :param name: A name for the model run.
@@ -83,6 +84,8 @@ def create_model_run(
         :param model_id: The internally-controlled identifier of the model.
                     The 'reference_id' field should be empty if this field is populated,
         :param metadata: An arbitrary metadata blob for the current run.
+        :param annotation_metadata_schema: A dictionary that defines schema for annotations.
+        :param segmentation_metadata_schema: A dictionary that defines schema for segmentation.
 
         :return:
         {
@@ -91,7 +94,11 @@ def create_model_run(
         }
         """
         payload = construct_model_run_creation_payload(
-            name, reference_id, model_id, metadata
+            name,
+            reference_id,
+            model_id,
+            metadata,
+            annotation_metadata_schema,
         )
         return self._client.create_model_run(self.id, payload)
 

nucleus/payload_constructor.py

@@ -19,6 +19,7 @@
     ITEMS_KEY,
     FORCE_KEY,
     MODEL_ID_KEY,
+    ANNOTATION_METADATA_SCHEMA_KEY,
     SEGMENTATIONS_KEY,
 )
 
@@ -87,6 +88,7 @@ def construct_model_run_creation_payload(
     reference_id: Optional[str],
     model_id: Optional[str],
     metadata: Optional[Dict],
+    annotation_metadata_schema: Optional[Dict] = None,
 ) -> dict:
     payload = {
         NAME_KEY: name,
@@ -100,4 +102,5 @@ def construct_model_run_creation_payload(
         NAME_KEY: name,
         REFERENCE_ID_KEY: reference_id,
         METADATA_KEY: metadata if metadata else {},
+        ANNOTATION_METADATA_SCHEMA_KEY: annotation_metadata_schema,
     }

nucleus/utils.py

@@ -0,0 +1,36 @@
+from typing import List, Union, Dict
+
+from .dataset_item import DatasetItem
+from .prediction import BoxPrediction, PolygonPrediction
+
+
+def _get_all_field_values(metadata_list: List[dict], key: str):
+    return {metadata[key] for metadata in metadata_list if key in metadata}
+
+
+def suggest_metadata_schema(
+    data: Union[
+        List[DatasetItem], List[BoxPrediction], List[PolygonPrediction]
+    ]
+):
+    metadata_list: List[dict] = [
+        d.metadata for d in data if d.metadata is not None
+    ]
+    schema = {}
+    all_keys = {k for metadata in metadata_list for k in metadata.keys()}
+
+    all_key_values: Dict[str, set] = {
+        k: _get_all_field_values(metadata_list, k) for k in all_keys
+    }
+
+    for key, values in all_key_values.items():
+        entry: dict = {}
+        if all(isinstance(x, (float, int)) for x in values):
+            entry["type"] = "number"
+        elif len(values) <= 50:
+            entry["type"] = "category"
+            entry["choices"] = list(values)
+        else:
+            entry["type"] = "text"
+        schema[key] = entry
+    return schema