Export category selection

Author: jsbroks

backend/database/datasets.py

@@ -69,18 +69,21 @@ def import_coco(self, coco_json):
             "name": task.name
         }
 
-    def export_coco(self, style="COCO"):
+    def export_coco(self, categories=None, style="COCO"):
 
         from workers.tasks import export_annotations
 
+        if categories is None or len(categories) == 0:
+            categories = self.categories
+
         task = TaskModel(
             name=f"Exporting {self.name} into {style} format",
             dataset_id=self.id,
             group="Annotation Export"
         )
         task.save()
 
-        cel_task = export_annotations.delay(task.id, self.id)
+        cel_task = export_annotations.delay(task.id, self.id, categories)
 
         return {
             "celery_id": cel_task.id,

backend/database/exports.py

@@ -10,6 +10,7 @@ class ExportModel(DynamicDocument):
     dataset_id = IntField(required=True)
     path = StringField(required=True)
     tags = ListField(default=[])
+    categories = ListField(default=[])
     created_at = DateTimeField(default=datetime.datetime.utcnow)
 
     def get_file(self):

backend/webserver/api/datasets.py

@@ -42,6 +42,8 @@
 coco_upload = reqparse.RequestParser()
 coco_upload.add_argument('coco', location='files', type=FileStorage, required=True, help='Json coco')
 
+export = reqparse.RequestParser()
+export.add_argument('categories', type=str, default=None, required=False, help='Ids of categories to export')
 
 update_dataset = reqparse.RequestParser()
 update_dataset.add_argument('categories', location='json', type=list, help="New list of categories")
@@ -173,6 +175,7 @@ def get(self, dataset_id):
 
 @api.route('/<int:dataset_id>')
 class DatasetId(Resource):
+
     @login_required
     def delete(self, dataset_id):
         """ Deletes dataset by ID (only owners)"""
@@ -341,13 +344,16 @@ def get(self, dataset_id):
 
         subdirectories = [f for f in sorted(os.listdir(directory))
                           if os.path.isdir(directory + f) and not f.startswith('.')]
+        
+        categories = CategoryModel.objects(id__in=dataset.categories).only('id', 'name')
 
         return {
             "pagination": pagination.export(),
             "images": images_json,
             "folder": folder,
             "directory": directory,
             "dataset": query_util.fix_ids(dataset),
+            "categories": query_util.fix_ids(categories),
             "subdirectories": subdirectories
         }
 
@@ -384,6 +390,26 @@ def get(self, dataset_id):
 @api.route('/<int:dataset_id>/export')
 class DatasetExport(Resource):
 
+    @api.expect(export)
+    @login_required
+    def get(self, dataset_id):
+
+        args = export.parse_args()
+        categories = args.get('categories')
+        
+        if len(categories) == 0:
+            categories = []
+
+        if len(categories) > 0 or isinstance(categories, str):
+            categories = [int(c) for c in categories.split(',')]
+
+        dataset = DatasetModel.objects(id=dataset_id).first()
+        
+        if not dataset:
+            return {'message': 'Invalid dataset ID'}, 400
+        
+        return dataset.export_coco(categories=categories)
+    
     @api.expect(coco_upload)
     @login_required
     def post(self, dataset_id):
@@ -412,7 +438,7 @@ def get(self, dataset_id):
         if not current_user.can_download(dataset):
             return {"message": "You do not have permission to download the dataset's annotations"}, 403
 
-        return coco_util.get_dataset_coco(dataset)
+        return coco_util.get_dataseext_coco(dataset)
 
     @api.expect(coco_upload)
     @login_required
@@ -459,16 +485,3 @@ def get(self, dataset_id):
 
         return dataset.scan()
 
-
-@api.route('/<int:dataset_id>/export')
-class DatasetScan(Resource):
-    
-    @login_required
-    def get(self, dataset_id):
-
-        dataset = DatasetModel.objects(id=dataset_id).first()
-        
-        if not dataset:
-            return {'message': 'Invalid dataset ID'}, 400
-        
-        return dataset.export_coco()

backend/workers/tasks/data.py

@@ -20,7 +20,7 @@
 
 
 @shared_task
-def export_annotations(task_id, dataset_id):
+def export_annotations(task_id, dataset_id, categories):
 
     task = TaskModel.objects.get(id=task_id)
     dataset = DatasetModel.objects.get(id=dataset_id)
@@ -30,11 +30,11 @@ def export_annotations(task_id, dataset_id):
 
     task.info("Beginning Export (COCO Format)")
 
-    db_categories = CategoryModel.objects(id__in=dataset.categories, deleted=False) \
+    db_categories = CategoryModel.objects(id__in=categories, deleted=False) \
         .only(*CategoryModel.COCO_PROPERTIES)
     db_images = ImageModel.objects(deleted=False, annotated=True, dataset_id=dataset.id)\
         .only(*ImageModel.COCO_PROPERTIES)
-    db_annotations = AnnotationModel.objects(deleted=False)
+    db_annotations = AnnotationModel.objects(deleted=False, category_id__in=categories)
 
     total_items = db_categories.count()
 
@@ -48,6 +48,7 @@ def export_annotations(task_id, dataset_id):
     progress = 0
 
     # iterate though all categoires and upsert
+    category_names = []
     for category in fix_ids(db_categories):
 
         if len(category.get('keypoint_labels', [])) > 0:
@@ -61,6 +62,7 @@ def export_annotations(task_id, dataset_id):
 
         task.info(f"Adding category: {category.get('name')}")
         coco.get('categories').append(category)
+        category_names.append(category.get('name'))
 
         progress += 1
         task.set_progress((progress/total_items)*100, socket=socket)
@@ -111,7 +113,7 @@ def export_annotations(task_id, dataset_id):
         json.dump(coco, fp)
 
     task.info("Creating export object")
-    export = ExportModel(dataset_id=dataset.id, path=file_path, tags=["COCO"])
+    export = ExportModel(dataset_id=dataset.id, path=file_path, tags=["COCO", *category_names])
     export.save()
 
     task.set_progress(100, socket=socket)

client/src/components/Pagination.vue

@@ -95,7 +95,7 @@ export default {
 
 <style>
 .page {
-    display:block;
-    margin:0 auto;
+  display: block;
+  margin: 0 auto;
 }
 </style>

client/src/models/datasets.js

@@ -30,8 +30,8 @@ export default {
   scan(id) {
     return axios.get(`${baseURL}/${id}/scan`);
   },
-  exportingCOCO(id) {
-    return axios.get(`${baseURL}/${id}/export`);
+  exportingCOCO(id, categories) {
+    return axios.get(`${baseURL}/${id}/export?categories=${categories}`);
   },
   getCoco(id) {
     return axios.get(`${baseURL}/${id}/coco`);