Merge pull request #39 from IDEA-Research/feature/export_label_project

Feature/export label project

deepdataspace/constants.py

@@ -261,6 +261,7 @@ class LabelProjectStatus:
     Reviewing = "reviewing"  #: Manager has finished the label project, waiting for owner to review.
     Rejected = "rejected"  #: Owner rejected the project.
     Accepted = "accepted"  #: Owner accepted the project.
+    Exported = "exported"  #: Owner has exported the project back to datasets.
 
 
 class LabelProjectRoles:

deepdataspace/model/image.py

@@ -265,6 +265,35 @@ def _update_dataset(self, bbox, segmentation, alpha_uri, coco_keypoints):
         if modified:
             self.dataset.save()
 
+    def _add_annotation(self,
+                        category: str,
+                        label: str = LabelName.GroundTruth,
+                        label_type: Literal["GT", "Pred", "User"] = "GT",
+                        conf: float = 1.0,
+                        is_group: bool = False,
+                        bbox: Tuple[int, int, int, int] = None,
+                        segmentation: List[List[int]] = None,
+                        alpha_uri: str = None,
+                        coco_keypoints: List[Union[float, int]] = None,
+                        confirm_type: int = 0,
+                        ):
+        if bbox:
+            if not self.width or not self.height:
+                raise ValueError("image width and height must be set before setting bbox")
+
+        label_obj = self._add_label(label, label_type)
+        category_obj = self._add_category(category)
+        bounding_box = self._format_bbox(self.width, self.height, bbox)
+        segmentation = self._format_segmentation(segmentation)
+        points, colors, lines, names = self._format_coco_keypoints(coco_keypoints)
+        anno_obj = Object(label_name=label, label_type=label_type, label_id=label_obj.id,
+                          category_name=category, category_id=category_obj.id,
+                          bounding_box=bounding_box, segmentation=segmentation, alpha=alpha_uri,
+                          points=points, lines=lines, point_colors=colors, point_names=names,
+                          conf=conf, is_group=is_group, confirm_type=confirm_type)
+        anno_obj.post_init()
+        self.objects.append(anno_obj)
+
     def add_annotation(self,
                        category: str,
                        label: str = LabelName.GroundTruth,
@@ -294,22 +323,9 @@ def add_annotation(self,
         :param confirm_type: the confirm_type of the annotation, 0 = not confirmed, 1 = gt may be fn, 2 = pred may be fp
         """
 
-        if bbox:
-            if not self.width or not self.height:
-                raise ValueError("image width and height must be set before setting bbox")
-
-        label_obj = self._add_label(label, label_type)
-        category_obj = self._add_category(category)
-        bounding_box = self._format_bbox(self.width, self.height, bbox)
-        segmentation = self._format_segmentation(segmentation)
-        points, colors, lines, names = self._format_coco_keypoints(coco_keypoints)
-        anno_obj = Object(label_name=label, label_type=label_type, label_id=label_obj.id,
-                          category_name=category, category_id=category_obj.id,
-                          bounding_box=bounding_box, segmentation=segmentation, alpha=alpha_uri,
-                          points=points, lines=lines, point_colors=colors, point_names=names,
-                          conf=conf, is_group=is_group, confirm_type=confirm_type)
-        anno_obj.post_init()
-        self.objects.append(anno_obj)
+        self._add_annotation(category, label, label_type, conf,
+                             is_group, bbox, segmentation, alpha_uri,
+                             coco_keypoints, confirm_type)
 
         self.save()
         self._update_dataset(bbox, segmentation, alpha_uri, coco_keypoints)

deepdataspace/model/label_task.py

@@ -5,6 +5,7 @@
 """
 
 import copy
+import logging
 import time
 import uuid
 from typing import ClassVar
@@ -18,20 +19,30 @@
 from pymongo.collection import Collection
 from pymongo.typings import _DocumentType
 
+from deepdataspace.constants import AnnotationType
+from deepdataspace.constants import DatasetStatus
 from deepdataspace.constants import LabelImageQAActions
 from deepdataspace.constants import LabelProjectQAActions
 from deepdataspace.constants import LabelProjectRoles
 from deepdataspace.constants import LabelProjectStatus
 from deepdataspace.constants import LabelTaskImageStatus
 from deepdataspace.constants import LabelTaskQAActions
 from deepdataspace.constants import LabelTaskStatus
+from deepdataspace.constants import LabelType
 from deepdataspace.model._base import BaseModel
+from deepdataspace.model.category import Category
 from deepdataspace.model.dataset import DataSet
 from deepdataspace.model.image import Image
+from deepdataspace.model.image import ImageModel
+from deepdataspace.model.label import Label
+from deepdataspace.model.object import Object
 from deepdataspace.model.user import User
+from deepdataspace.utils.string import get_str_md5
 
 Num = Union[float, int]
 
+logger = logging.getLogger("io.model.label_task")
+
 
 def current_ts():
     return int(time.time() * 1000)
@@ -378,6 +389,143 @@ def qa_project(self, action):
         else:
             raise LabelProjectError(f"Invalid project qa action: {action}.")
 
+    @staticmethod
+    def _get_image_batch(dataset_id, offset):
+        """
+        Get a batch of images from the dataset.
+        """
+
+        IModel = Image(dataset_id)
+        images: Dict[int, ImageModel] = {i.id: i for i in
+                                         IModel.find_many({}, sort=[("id", 1)], size=100, skip=offset)}
+        return images, offset + 100
+
+    @staticmethod
+    def _get_label(dataset_id: str, label_set_name: str):
+        """
+        Get the label set of saving annotations.
+        """
+        label_id = get_str_md5(f"{dataset_id}_{label_set_name}")
+        label_obj = Label(name=label_set_name, id=label_id, type=LabelType.GroundTruth, dataset_id=dataset_id)
+        label_obj.post_init()
+        label_obj.save()
+        return label_obj
+
+    @staticmethod
+    def _get_category(dataset_id: str, category_name: str, categories: dict):
+        """
+        Get the category of saving annotation.
+        """
+        cat_obj = categories.get(category_name, None)
+        if cat_obj is None:
+            cat_id = get_str_md5(f"{dataset_id}_{category_name}")
+            cat_obj = Category(id=cat_id, name=category_name, dataset_id=dataset_id)
+            cat_obj.post_init()
+            cat_obj.save()
+            categories[category_name] = cat_obj
+        return cat_obj
+
+    def _export_dataset(self, dataset: DataSet, label_set_name: str):
+        # the label set
+        label_obj = self._get_label(dataset.id, label_set_name)
+
+        # the categories cache, we cache it to avoid duplicated db query and insertion
+        categories = {}
+
+        # the queue of target images
+        LTImage = LabelTaskImage(dataset.id)
+        images, offset = self._get_image_batch(dataset.id, 0)
+
+        has_bbox = False  # whether the annotations have bbox
+
+        # iter every label image, save every annotation to target image
+        for ltimage in LTImage.find_many({}, sort=[("image_id", 1)]):
+            if not images:  # no more images in the queue, get a new batch
+                images, offset = self._get_image_batch(dataset.id, offset)
+
+            # match label image and target image
+            image_id = ltimage.image_id
+            image = images.pop(image_id, None)
+            if image is None:
+                continue
+            image.objects = [o for o in image.objects if o.label_id != label_obj.id]
+
+            # save annotations
+            labels = ltimage.labels
+            for labeler_id, label_list in labels.items():
+                if not label_list:
+                    continue
+
+                label_data = label_list[0]
+                annotations = label_data.annotations
+
+                for anno in annotations:
+                    category = anno["category_name"]
+                    if not category:
+                        continue
+
+                    bounding_box = anno["bounding_box"]
+                    if not bounding_box:
+                        continue
+
+                    has_bbox = True
+                    cat_obj = self._get_category(dataset.id, category, categories)
+
+                    anno_obj = Object(label_name=label_obj.name, label_type=label_obj.type, label_id=label_obj.id,
+                                      category_name=cat_obj.name, category_id=cat_obj.id,
+                                      bounding_box=anno["bounding_box"])
+                    anno_obj.post_init()
+                    image.objects.append(anno_obj)
+                    image.batch_save()
+
+        Image(dataset.id).finish_batch_save()
+
+        if has_bbox:
+            if AnnotationType.Classification not in dataset.object_types:
+                dataset.object_types.append(AnnotationType.Classification)
+            if AnnotationType.Detection not in dataset.object_types:
+                dataset.object_types.append(AnnotationType.Detection)
+            dataset.save()
+
+    def export_project(self, label_set_name: str):
+        """
+        Export the label data back to datasets.
+        """
+
+        if self.status != LabelProjectStatus.Accepted:
+            msg = f"Project can only be exported in status of 'accepted', current status is {self.status}."
+            raise LabelProjectError(msg)
+
+        label_set_name = f"{label_set_name}[{self.id[:8]}]"
+        logger.info(f"exporting project {self.id} to label set {label_set_name}")
+
+        num = len(self.datasets)
+        for idx, dataset in enumerate(self.datasets):
+            dataset_id = dataset["id"]
+            dataset = DataSet.find_one({"id": dataset_id})
+            if dataset is None:
+                continue
+
+            status = DatasetStatus.Ready
+            try:
+                update_data = {"status"                            : DatasetStatus.Importing,
+                               "detail_status.export_label_project": DatasetStatus.Importing}
+                DataSet.update_one({"id": dataset_id}, update_data)
+                logger.info(f"[{idx + 1}/{num}]exporting label project to dataset {dataset_id}")
+                self._export_dataset(dataset, label_set_name)
+                self.status = LabelProjectStatus.Exported
+                self.save()
+            except Exception as e:
+                status = DatasetStatus.Failed
+                logger.warning(f"[{idx + 1}/{num}]export label project to dataset {dataset_id} failed: {e}")
+            else:
+                status = DatasetStatus.Ready
+                logger.info(f"[{idx + 1}/{num}]export label project to dataset {dataset_id} success")
+            finally:
+                update_data = {"status"                            : DatasetStatus.Ready,
+                               "detail_status.export_label_project": status}
+                DataSet.update_one({"id": dataset_id}, update_data)
+
 
 class ProjectRole(BaseModel):
     """
@@ -633,6 +781,14 @@ def can_qa_project(user: User, project_id):
 
         return ProjectRole.is_owner(user, project_id)
 
+    @staticmethod
+    def can_export_project(user: User, project_id):
+        """
+        Check if target user can export the project.
+        """
+
+        return ProjectRole.is_owner(user, project_id)
+
 
 class TaskRole(BaseModel):
     """

deepdataspace/server/resources/api_v1/__init__.py

@@ -36,6 +36,7 @@
     path("label_projects/<project_id>", label_tasks.ProjectView.as_view()),
     path("label_project_configs/<project_id>", label_tasks.ProjectConfigView.as_view()),
     path("label_project_qa/<project_id>", label_tasks.ProjectQAView.as_view()),
+    path("label_project_export/<project_id>", label_tasks.ProjectExportView.as_view()),
     path("label_tasks", label_tasks.TasksView.as_view()),
     path("label_task_configs/<task_id>", label_tasks.TaskConfigView.as_view()),
     path("label_task_roles/<task_id>", label_tasks.TaskRolesView.as_view()),

deepdataspace/server/resources/api_v1/label_tasks.py

@@ -323,6 +323,40 @@ def post(self, request, project_id):
         return format_response({})
 
 
+class ProjectExportView(AuthenticatedAPIView):
+    """
+    - POST /api/v1/label_project_export/<project_id>
+    """
+
+    post_args = [
+        Argument("label_name", str, Argument.JSON, required=True),
+    ]
+
+    def post(self, request, project_id):
+        """
+        Export a label project back to datasets.
+
+        - POST /api/v1/label_project_export/<project_id>
+        """
+
+        user = request.user
+        label_name, = parse_arguments(request, self.post_args)
+
+        project = LabelProject.find_one({"id": project_id})
+        if project is None:
+            raise_exception(404, f"project[id={project_id}] is not found")
+
+        if not ProjectRole.can_export_project(user, project_id):
+            raise_exception(403, f"user[{user.id}] is not allowed to export project[{project_id}]")
+
+        try:
+            project.export_project(label_name)
+        except LabelProjectError as err:
+            raise_exception(400, str(err))
+
+        return format_response({})
+
+
 class TasksView(AuthenticatedAPIView):
     """
     - GET /api/v1/label_tasks

packages/app/src/locales/en-US.ts

@@ -381,10 +381,12 @@ export default {
   'proj.infoModal.label': 'Project Managers',
 
   'proj.exportModal.title': 'Export to Dataset',
-  'proj.exportModal.labelSet.name': 'Labelset',
-  'proj.exportModal.labelSet.rule': 'Please input labelset name',
+  'proj.exportModal.labelName.name': 'Labelset',
+  'proj.exportModal.labelName.rule': 'Please input labelset name',
+  'proj.exportModal.labelName.tips':
+    'You can check labeling result in selected dataset with the labelset name after clicking "OK". ',
   'proj.exportModal.submitSuccess':
-    'Successfully export to labelset "{name}", you can check results in datasets module.',
+    'Successfully export labelset "{name}" into selected dataset, you can check labeled results in datasets module.',
 
   'proj.workspace.eTask.startLabel': 'Start Label',
   'proj.workspace.eTask.edit': 'Edit',
@@ -401,6 +403,7 @@ export default {
   'proj.statusMap.reviewing': 'Reviewing',
   'proj.statusMap.rejected': 'Rejected',
   'proj.statusMap.accepted': 'Accepted',
+  'proj.statusMap.exported': 'Exported',
 
   'proj.eTaskStatus.waiting': 'Waiting',
   'proj.eTaskStatus.working': 'Working',

packages/app/src/locales/zh-CN.ts

@@ -350,10 +350,12 @@ export default {
   'proj.infoModal.label': '项目经理',
 
   'proj.exportModal.title': '导出到数据集',
-  'proj.exportModal.labelSet.name': '标注集名称',
-  'proj.exportModal.labelSet.rule': '请输入标注集名称',
+  'proj.exportModal.labelName.name': '标注集名称',
+  'proj.exportModal.labelName.rule': '请输入标注集名称',
+  'proj.exportModal.labelName.tips':
+    '点击“确定”后，可以用标注集名称查看所选数据集的标注结果。',
   'proj.exportModal.submitSuccess':
-    '已成功导入至 "{name}" 结果集中，您可以移步数据集模块查看标注结果',
+    '已成功导出标注集 "{name}" 到所选数据集，您可以在数据集模块中查看标注结果。',
 
   'proj.workspace.eTask.startLabel': '开始标注',
   'proj.workspace.eTask.edit': '编辑',
@@ -370,6 +372,7 @@ export default {
   'proj.statusMap.reviewing': '审核中',
   'proj.statusMap.rejected': '已拒绝',
   'proj.statusMap.accepted': '已通过',
+  'proj.statusMap.exported': '已导出',
 
   'proj.eTaskStatus.waiting': '等待中',
   'proj.eTaskStatus.working': '进行中',