diff --git a/backend/database/__init__.py b/backend/database/__init__.py index 8f1ab9fb..046c4661 100644 --- a/backend/database/__init__.py +++ b/backend/database/__init__.py @@ -40,9 +40,9 @@ def upsert(model, query=None, update=None): return new_model -def fix_ids(objs): - objects_list = json.loads(objs.to_json().replace('\"_id\"', '\"id\"')) - return objects_list +def fix_ids(q): + json_obj = json.loads(q.to_json().replace('\"_id\"', '\"id\"')) + return json_obj def create_from_json(json_file): diff --git a/backend/database/annotations.py b/backend/database/annotations.py index ccdb6dcf..b0c230a8 100644 --- a/backend/database/annotations.py +++ b/backend/database/annotations.py @@ -10,6 +10,10 @@ class AnnotationModel(DynamicDocument): + COCO_PROPERTIES = ["id", "image_id", "category_id", "segmentation", \ + "iscrowd", "color", "area", "bbox", "metadata", \ + "keypoints"] + id = SequenceField(primary_key=True) image_id = IntField(required=True) category_id = IntField(required=True) @@ -41,13 +45,15 @@ class AnnotationModel(DynamicDocument): def __init__(self, image_id=None, **data): from .images import ImageModel - image = ImageModel.objects(id=image_id).first() - if image is not None: - data['image_id'] = image_id - data['width'] = image.width - data['height'] = image.height - data['dataset_id'] = image.dataset_id + if image_id is not None: + image = ImageModel.objects(id=image_id).first() + + if image is not None: + data['image_id'] = image_id + data['width'] = image.width + data['height'] = image.height + data['dataset_id'] = image.dataset_id super(AnnotationModel, self).__init__(**data) diff --git a/backend/database/categories.py b/backend/database/categories.py index 5dcdd530..49074d69 100644 --- a/backend/database/categories.py +++ b/backend/database/categories.py @@ -7,6 +7,9 @@ class CategoryModel(DynamicDocument): + COCO_PROPERTIES = ["id", "name", "supercategory", "color", "metadata",\ + "keypoint_edges", "keypoint_labels"] + id = SequenceField(primary_key=True) name = StringField(required=True, unique_with=['creator']) supercategory = StringField(default='') diff --git a/backend/database/images.py b/backend/database/images.py index 8a210111..db7b88a1 100644 --- a/backend/database/images.py +++ b/backend/database/images.py @@ -10,7 +10,10 @@ from .annotations import AnnotationModel class ImageModel(DynamicDocument): - + + COCO_PROPERTIES = ["id", "width", "height", "file_name", "path", "license",\ + "flickr_url", "coco_url", "date_captured", "dataset_id"] + # -- Contants THUMBNAIL_DIRECTORY = '.thumbnail' PATTERN = (".gif", ".png", ".jpg", ".jpeg", ".bmp") @@ -37,6 +40,7 @@ class ImageModel(DynamicDocument): thumbnail_url = StringField() image_url = StringField() coco_url = StringField() + date_captured = DateTimeField() metadata = DictField() license = IntField() diff --git a/backend/webserver/api/images.py b/backend/webserver/api/images.py index 05ea75e7..3b232523 100644 --- a/backend/webserver/api/images.py +++ b/backend/webserver/api/images.py @@ -204,5 +204,5 @@ def get(self, image_id): if not current_user.can_download(image): return {"message": "You do not have permission to download the images's annotations"}, 403 - return coco_util.get_image_coco(image) + return coco_util.get_image_coco(image_id) diff --git a/backend/webserver/util/coco_util.py b/backend/webserver/util/coco_util.py index 9d520e68..0a43e5d6 100644 --- a/backend/webserver/util/coco_util.py +++ b/backend/webserver/util/coco_util.py @@ -1,7 +1,7 @@ import pycocotools.mask as mask -from .query_util import fix_ids from database import ( + fix_ids, ImageModel, DatasetModel, CategoryModel, @@ -102,52 +102,52 @@ def get_annotations_iou(annotation_a, annotation_b): return ious[0][0] -def get_image_coco(image): +def get_image_coco(image_id): """ Generates coco for an image :param image: ImageModel :return: Coco in dictionary format """ - dataset = DatasetModel.objects(id=image.dataset_id).first() - image = fix_ids(image) + image = ImageModel.objects(id=image_id)\ + .only(*ImageModel.COCO_PROPERTIES) + + image = fix_ids(image)[0] + dataset = DatasetModel.objects(id=image.get('dataset_id')).first() - bulk_categories = CategoryModel.objects(deleted=False) \ - .exclude('deleted_date').in_bulk(dataset.categories).items() + bulk_categories = CategoryModel.objects(id__in=dataset.categories, deleted=False) \ + .only(*CategoryModel.COCO_PROPERTIES) + print(bulk_categories) + + db_annotations = AnnotationModel.objects(deleted=False, image_id=image_id) categories = [] annotations = [] - for category in bulk_categories: - category = category[1] - category_annotations = AnnotationModel.objects( - deleted=False, category_id=category.id, image_id=image.get('id') - ).exclude('paper_object', 'deleted_date').all() + for category in fix_ids(bulk_categories): + + category_annotations = db_annotations\ + .filter(category_id=category.get('id'))\ + .only(*AnnotationModel.COCO_PROPERTIES) - if len(category_annotations) == 0: + if category_annotations.count() == 0: continue + category_annotations = fix_ids(category_annotations) for annotation in category_annotations: - annotation = fix_ids(annotation) has_segmentation = len(annotation.get('segmentation', [])) > 0 has_keypoints = len(annotation.get('keypoints', [])) > 0 if has_segmentation or has_keypoints: - del annotation['deleted'] - if not has_keypoints: - del annotation['keypoints'] - else: + if has_keypoints: arr = np.array(annotation.get('keypoints', [])) arr = arr[2::3] annotation['num_keypoints'] = len(arr[arr > 0]) - + annotations.append(annotation) - category = fix_ids(category) - del category['deleted'] - if len(category.get('keypoint_labels')) > 0: category['keypoints'] = category.pop('keypoint_labels') category['skeleton'] = category.pop('keypoint_edges') @@ -157,8 +157,6 @@ def get_image_coco(image): categories.append(category) - del image['deleted'] - coco = { "images": [image], "categories": categories, diff --git a/backend/workers/tasks/data.py b/backend/workers/tasks/data.py index ea42988c..e599fe5c 100644 --- a/backend/workers/tasks/data.py +++ b/backend/workers/tasks/data.py @@ -27,34 +27,34 @@ def export_annotations(task_id, dataset_id): task.info("Beginning Export (COCO Format)") - categories = CategoryModel.objects(deleted=False) \ - .exclude('deleted_date').in_bulk(dataset.categories).items() + db_categories = CategoryModel.objects(id__in=dataset.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) - total_items = len(dataset.categories) + total_items = db_categories.count() dataset = fix_ids(dataset) - images = ImageModel.objects(deleted=False, annotated=True, dataset_id=dataset.get('id')).exclude('deleted_date') - all_annotations = AnnotationModel.objects(deleted=False).exclude('deleted_date', 'paper_object') - coco = { 'images': [], 'categories': [], 'annotations': [] } - total_items += images.count() + total_items += db_images.count() progress = 0 - for category in categories: - category = fix_ids(category[1]) - del category['deleted'] + # iterate though all categoires and upsert + for category in fix_ids(db_categories): + if len(category.get('keypoint_labels', [])) > 0: category['keypoints'] = category.pop('keypoint_labels', []) category['skeleton'] = category.pop('keypoint_edges', []) else: if 'keypoint_edges' in category: del category['keypoint_edges'] - if 'keypoint_labels' in categories: + if 'keypoint_labels' in category: del category['keypoint_labels'] task.info(f"Adding category: {category.get('name')}") @@ -63,14 +63,16 @@ def export_annotations(task_id, dataset_id): progress += 1 task.set_progress((progress/total_items)*100, socket=socket) - total_annotations = 0 - total_images = 0 - for image in images: - annotations = all_annotations.filter(image_id=image.id) - if annotations.count() == 0: - continue + total_annotations = db_annotations.count() + total_images = db_images.count() + for image in fix_ids(db_images): + + progress += 1 + task.set_progress((progress/total_items)*100, socket=socket) - annotations = fix_ids(annotations.all()) + annotations = db_annotations.filter(image_id=image.get('id'))\ + .only(*AnnotationModel.COCO_PROPERTIES) + annotations = fix_ids(annotations) num_annotations = 0 for annotation in annotations: @@ -78,7 +80,6 @@ def export_annotations(task_id, dataset_id): has_segmentation = len(annotation.get('segmentation', [])) > 0 if has_keypoints or has_segmentation: - del annotation['deleted'] if not has_keypoints: if 'keypoints' in annotation: @@ -91,16 +92,8 @@ def export_annotations(task_id, dataset_id): num_annotations += 1 coco.get('annotations').append(annotation) - task.info(f'Exporting {num_annotations} annotations for image {image.id}') - total_annotations += num_annotations - - image = fix_ids(image) - del image['deleted'] + task.info(f"Exporting {num_annotations} annotations for image {image.get('id')}") coco.get('images').append(image) - total_images += 1 - - progress += 1 - task.set_progress((progress/total_items)*100, socket=socket) file_path = dataset.get('directory') + '/coco.json' with open(file_path, 'w') as fp: