Spaces:
Runtime error
Runtime error
File size: 2,241 Bytes
989ec3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from flask_login import current_user
from mongoengine import *
import imantics as im
class CategoryModel(DynamicDocument):
COCO_PROPERTIES = ["id", "name", "supercategory", "color", "metadata",\
"keypoint_edges", "keypoint_labels", "keypoint_colors"]
id = SequenceField(primary_key=True)
name = StringField(required=True, unique_with=['creator'])
supercategory = StringField(default='')
color = StringField(default=None)
metadata = DictField(default={})
creator = StringField(default='unknown')
deleted = BooleanField(default=False)
deleted_date = DateTimeField()
keypoint_edges = ListField(default=[])
keypoint_labels = ListField(default=[])
keypoint_colors = ListField(default=[])
@classmethod
def bulk_create(cls, categories):
if not categories:
return []
category_ids = []
for category in categories:
category_model = CategoryModel.objects(name=category).first()
if category_model is None:
new_category = CategoryModel(name=category)
new_category.save()
category_ids.append(new_category.id)
else:
category_ids.append(category_model.id)
return category_ids
def save(self, *args, **kwargs):
if not self.color:
self.color = im.Color.random().hex
if current_user:
self.creator = current_user.username
else:
self.creator = 'system'
return super(CategoryModel, self).save(*args, **kwargs)
def __call__(self):
""" Generates imantics category object """
data = {
'name': self.name,
'color': self.color,
'parent': self.supercategory,
'metadata': self.metadata,
'id': self.id
}
return im.Category(**data)
def is_owner(self, user):
if user.is_admin:
return True
return user.username.lower() == self.creator.lower()
def can_edit(self, user):
return self.is_owner(user)
def can_delete(self, user):
return self.is_owner(user)
__all__ = ["CategoryModel"] |