Spaces:
Runtime error
Runtime error
File size: 3,337 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
import imantics as im
import json
from mongoengine import *
from .datasets import DatasetModel
from .categories import CategoryModel
from .events import Event
from flask_login import current_user
import numpy as np
import cv2
class AnnotationModel(DynamicDocument):
COCO_PROPERTIES = ["id", "image_id", "category_id", "segmentation",
"iscrowd", "color", "area", "bbox", "metadata",
"keypoints", "isbbox"]
id = SequenceField(primary_key=True)
image_id = IntField(required=True)
category_id = IntField(required=True)
dataset_id = IntField()
segmentation = ListField(default=[])
area = IntField(default=0)
bbox = ListField(default=[0, 0, 0, 0])
iscrowd = BooleanField(default=False)
isbbox = BooleanField(default=False)
creator = StringField(required=True)
width = IntField()
height = IntField()
color = StringField()
keypoints = ListField(default=[])
metadata = DictField(default={})
paper_object = ListField(default=[])
deleted = BooleanField(default=False)
deleted_date = DateTimeField()
milliseconds = IntField(default=0)
events = EmbeddedDocumentListField(Event)
def __init__(self, image_id=None, **data):
from .images import ImageModel
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)
def save(self, copy=False, *args, **kwargs):
if self.dataset_id and not copy:
dataset = DatasetModel.objects(id=self.dataset_id).first()
if dataset is not None:
self.metadata = dataset.default_annotation_metadata.copy()
if self.color is None:
self.color = im.Color.random().hex
if current_user:
self.creator = current_user.username
else:
self.creator = 'system'
return super(AnnotationModel, self).save(*args, **kwargs)
def is_empty(self):
return len(self.segmentation) == 0 or self.area == 0
def mask(self):
""" Returns binary mask of annotation """
mask = np.zeros((self.height, self.width))
pts = [
np.array(anno).reshape(-1, 2).round().astype(int)
for anno in self.segmentation
]
mask = cv2.fillPoly(mask, pts, 1)
return mask
def clone(self):
""" Creates a clone """
create = json.loads(self.to_json())
del create['_id']
return AnnotationModel(**create)
def __call__(self):
category = CategoryModel.objects(id=self.category_id).first()
if category:
category = category()
data = {
'image': None,
'category': category,
'color': self.color,
'polygons': self.segmentation,
'width': self.width,
'height': self.height,
'metadata': self.metadata
}
return im.Annotation(**data)
def add_event(self, e):
self.update(push__events=e)
__all__ = ["AnnotationModel"]
|