feat: filter labels when label_type_in is provided (#1783)

kili-technology · Oct 1, 2024 · 4d1e2fe · 4d1e2fe
1 parent 6a9b516
commit 4d1e2fe
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 2 deletions.
diff --git a/src/kili/services/export/tools.py b/src/kili/services/export/tools.py
@@ -165,6 +165,31 @@ def fetch_assets(
         kili.kili_api_gateway, download_media, fields, ProjectId(project_id), local_media_dir
     )
     assets_gen = kili.kili_api_gateway.list_assets(filters, fields, options)
+
+    if (label_type_in is not None) and (len(label_type_in) > 0):
+        if export_type == "latest":
+            assets_gen = filter(
+                lambda asset: asset["latestLabel"].get("labelType") in label_type_in, assets_gen
+            )
+        else:
+            assets_gen = filter(
+                lambda asset: any(
+                    label.get("labelType") in label_type_in for label in asset["labels"]
+                ),
+                assets_gen,
+            )
+            assets_gen = (
+                {
+                    **asset,
+                    "labels": [
+                        label
+                        for label in asset["labels"]
+                        if label.get("labelType") in label_type_in
+                    ],
+                }
+                for asset in assets_gen
+            )
+
     if download_media_function is not None:
         assets: List[Dict] = []
         # TODO: modify download_media function so it can take a generator of assets

diff --git a/tests/fakes/fake_data.py b/tests/fakes/fake_data.py
@@ -65,6 +65,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -77,6 +78,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection_with_classification,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -87,6 +89,7 @@
     "latestLabel": {
         "jsonResponse": {},
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_1",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_1.jpg",
@@ -97,6 +100,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_2",
     "content": "https://storage.googleapis.com/label-public-staging/car/car_2.jpg",
@@ -107,6 +111,7 @@
     "latestLabel": {
         "jsonResponse": job_object_detection,
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
+        "labelType": "DEFAULT",
     },
     "externalId": "car_3",
     "content": "",
@@ -134,6 +139,7 @@
             "1": job_object_detection,
             **{str(i): {} for i in range(2, 28)},
         },
+        "labelType": "DEFAULT",
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
     },
     "externalId": "short_video",
@@ -148,6 +154,7 @@
             "1": job_object_detection,
             **{str(i): {} for i in range(2, 130)},
         },
+        "labelType": "DEFAULT",
         "author": {"firstname": "Jean-Pierre", "lastname": "Dupont"},
     },
     "externalId": "video2",

diff --git a/tests/unit/adapters/kili_api_gateway/label/test_data/test_case_13.py b/tests/unit/adapters/kili_api_gateway/label/test_data/test_case_13.py
@@ -106,6 +106,7 @@
             "isLatestLabelForUser": True,
             "isSentBackToQueue": False,
             "jsonResponse": json_response,
+            "labelType": "DEFAULT",
         },
         "priority": 0,
         "skipped": False,

diff --git a/tests/unit/services/export/test_export.py b/tests/unit/services/export/test_export.py
@@ -826,7 +826,7 @@ def test_export_with_asset_filter_kwargs(mocker):
             "inference_mark_gte": 0.7,
             "inference_mark_lte": 0.8,
         },
-        label_type_in=["PREDICTION"],
+        label_type_in=["PREDICTION", "DEFAULT", "REVIEW"],
     )
     expected_where = AssetFilters(
         project_id=ProjectId("fake_proj_id"),
@@ -837,7 +837,7 @@ def test_export_with_asset_filter_kwargs(mocker):
         honeypot_mark_lte=0.4,
         created_at_gte="2022-02-03",
         created_at_lte="2023-02-03",
-        label_type_in=["PREDICTION"],
+        label_type_in=["PREDICTION", "DEFAULT", "REVIEW"],
         inference_mark_gte=0.7,
         inference_mark_lte=0.8,
         issue_type="QUESTION",