14
14
from cvat .apps .engine .models import Task , ShapeType , AttributeType
15
15
16
16
import datumaro .components .extractor as datumaro
17
- from datumaro .util .image import lazy_image
17
+ from datumaro .util .image import Image
18
18
19
19
20
20
class CvatImagesDirExtractor (datumaro .Extractor ):
@@ -29,8 +29,7 @@ def __init__(self, url):
29
29
path = osp .join (dirpath , name )
30
30
if self ._is_image (path ):
31
31
item_id = Task .get_image_frame (path )
32
- item = datumaro .DatasetItem (
33
- id = item_id , image = lazy_image (path ))
32
+ item = datumaro .DatasetItem (id = item_id , image = path )
34
33
items .append ((item .id , item ))
35
34
36
35
items = sorted (items , key = lambda e : int (e [0 ]))
@@ -49,112 +48,90 @@ def __len__(self):
49
48
def subsets (self ):
50
49
return self ._subsets
51
50
52
- def get (self , item_id , subset = None , path = None ):
53
- if path or subset :
54
- raise KeyError ()
55
- return self ._items [item_id ]
56
-
57
51
def _is_image (self , path ):
58
52
for ext in self ._SUPPORTED_FORMATS :
59
53
if osp .isfile (path ) and path .endswith (ext ):
60
54
return True
61
55
return False
62
56
63
57
64
- class CvatTaskExtractor (datumaro .Extractor ):
65
- def __init__ (self , url , db_task , user ):
66
- self ._db_task = db_task
67
- self ._categories = self ._load_categories ()
68
-
69
- cvat_annotations = TaskAnnotation (db_task .id , user )
70
- with transaction .atomic ():
71
- cvat_annotations .init_from_db ()
72
- cvat_annotations = Annotation (cvat_annotations .ir_data , db_task )
58
+ class CvatAnnotationsExtractor (datumaro .Extractor ):
59
+ def __init__ (self , url , cvat_annotations ):
60
+ self ._categories = self ._load_categories (cvat_annotations )
73
61
74
62
dm_annotations = []
75
63
76
- for cvat_anno in cvat_annotations .group_by_frame ():
77
- dm_anno = self ._read_cvat_anno (cvat_anno )
78
- dm_item = datumaro .DatasetItem (
79
- id = cvat_anno .frame , annotations = dm_anno )
64
+ for cvat_frame_anno in cvat_annotations .group_by_frame ():
65
+ dm_anno = self ._read_cvat_anno (cvat_frame_anno , cvat_annotations )
66
+ dm_image = Image (path = cvat_frame_anno .name , size = (
67
+ cvat_frame_anno .height , cvat_frame_anno .width )
68
+ )
69
+ dm_item = datumaro .DatasetItem (id = cvat_frame_anno .frame ,
70
+ annotations = dm_anno , image = dm_image )
80
71
dm_annotations .append ((dm_item .id , dm_item ))
81
72
82
73
dm_annotations = sorted (dm_annotations , key = lambda e : int (e [0 ]))
83
74
self ._items = OrderedDict (dm_annotations )
84
75
85
- self ._subsets = None
86
-
87
76
def __iter__ (self ):
88
77
for item in self ._items .values ():
89
78
yield item
90
79
91
80
def __len__ (self ):
92
81
return len (self ._items )
93
82
83
+ # pylint: disable=no-self-use
94
84
def subsets (self ):
95
- return self ._subsets
85
+ return []
86
+ # pylint: enable=no-self-use
96
87
97
- def get (self , item_id , subset = None , path = None ):
98
- if path or subset :
99
- raise KeyError ()
100
- return self ._items [item_id ]
88
+ def categories (self ):
89
+ return self ._categories
101
90
102
- def _load_categories (self ):
91
+ @staticmethod
92
+ def _load_categories (cvat_anno ):
103
93
categories = {}
104
94
label_categories = datumaro .LabelCategories ()
105
95
106
- db_labels = self ._db_task .label_set .all ()
107
- for db_label in db_labels :
108
- db_attributes = db_label .attributespec_set .all ()
109
- label_categories .add (db_label .name )
110
-
111
- for db_attr in db_attributes :
112
- label_categories .attributes .add (db_attr .name )
96
+ for _ , label in cvat_anno .meta ['task' ]['labels' ]:
97
+ label_categories .add (label ['name' ])
98
+ for _ , attr in label ['attributes' ]:
99
+ label_categories .attributes .add (attr ['name' ])
113
100
114
101
categories [datumaro .AnnotationType .label ] = label_categories
115
102
116
103
return categories
117
104
118
- def categories (self ):
119
- return self ._categories
120
-
121
- def _read_cvat_anno (self , cvat_anno ):
105
+ def _read_cvat_anno (self , cvat_frame_anno , cvat_task_anno ):
122
106
item_anno = []
123
107
124
108
categories = self .categories ()
125
109
label_cat = categories [datumaro .AnnotationType .label ]
126
-
127
- label_map = {}
128
- label_attrs = {}
129
- db_labels = self ._db_task .label_set .all ()
130
- for db_label in db_labels :
131
- label_map [db_label .name ] = label_cat .find (db_label .name )[0 ]
132
-
133
- attrs = {}
134
- db_attributes = db_label .attributespec_set .all ()
135
- for db_attr in db_attributes :
136
- attrs [db_attr .name ] = db_attr
137
- label_attrs [db_label .name ] = attrs
138
- map_label = lambda label_db_name : label_map [label_db_name ]
110
+ map_label = lambda name : label_cat .find (name )[0 ]
111
+ label_attrs = {
112
+ label ['name' ]: label ['attributes' ]
113
+ for _ , label in cvat_task_anno .meta ['task' ]['labels' ]
114
+ }
139
115
140
116
def convert_attrs (label , cvat_attrs ):
141
117
cvat_attrs = {a .name : a .value for a in cvat_attrs }
142
118
dm_attr = dict ()
143
- for attr_name , attr_spec in label_attrs [label ].items ():
144
- attr_value = cvat_attrs .get (attr_name , attr_spec .default_value )
119
+ for _ , a_desc in label_attrs [label ]:
120
+ a_name = a_desc ['name' ]
121
+ a_value = cvat_attrs .get (a_name , a_desc ['default_value' ])
145
122
try :
146
- if attr_spec . input_type == AttributeType .NUMBER :
147
- attr_value = float (attr_value )
148
- elif attr_spec . input_type == AttributeType .CHECKBOX :
149
- attr_value = attr_value .lower () == 'true'
150
- dm_attr [attr_name ] = attr_value
123
+ if a_desc [ ' input_type' ] == AttributeType .NUMBER :
124
+ a_value = float (a_value )
125
+ elif a_desc [ ' input_type' ] == AttributeType .CHECKBOX :
126
+ a_value = ( a_value .lower () == 'true' )
127
+ dm_attr [a_name ] = a_value
151
128
except Exception as e :
152
- slogger . task [ self . _db_task . id ]. error (
153
- "Failed to convert attribute '%s'='%s': %s" % \
154
- ( attr_name , attr_value , e ))
129
+ raise Exception (
130
+ "Failed to convert attribute '%s'='%s': %s" %
131
+ ( a_name , a_value , e ))
155
132
return dm_attr
156
133
157
- for tag_obj in cvat_anno .tags :
134
+ for tag_obj in cvat_frame_anno .tags :
158
135
anno_group = tag_obj .group
159
136
anno_label = map_label (tag_obj .label )
160
137
anno_attr = convert_attrs (tag_obj .label , tag_obj .attributes )
@@ -163,7 +140,7 @@ def convert_attrs(label, cvat_attrs):
163
140
attributes = anno_attr , group = anno_group )
164
141
item_anno .append (anno )
165
142
166
- for shape_obj in cvat_anno .labeled_shapes :
143
+ for shape_obj in cvat_frame_anno .labeled_shapes :
167
144
anno_group = shape_obj .group
168
145
anno_label = map_label (shape_obj .label )
169
146
anno_attr = convert_attrs (shape_obj .label , shape_obj .attributes )
@@ -183,8 +160,64 @@ def convert_attrs(label, cvat_attrs):
183
160
anno = datumaro .Bbox (x0 , y0 , x1 - x0 , y1 - y0 ,
184
161
label = anno_label , attributes = anno_attr , group = anno_group )
185
162
else :
186
- raise Exception ("Unknown shape type '%s'" % ( shape_obj .type ) )
163
+ raise Exception ("Unknown shape type '%s'" % shape_obj .type )
187
164
188
165
item_anno .append (anno )
189
166
190
- return item_anno
167
+ return item_anno
168
+
169
+
170
+ class CvatTaskExtractor (CvatAnnotationsExtractor ):
171
+ def __init__ (self , url , db_task , user ):
172
+ cvat_annotations = TaskAnnotation (db_task .id , user )
173
+ with transaction .atomic ():
174
+ cvat_annotations .init_from_db ()
175
+ cvat_annotations = Annotation (cvat_annotations .ir_data , db_task )
176
+ super ().__init__ (url , cvat_annotations )
177
+
178
+
179
+ def match_frame (item , cvat_task_anno ):
180
+ frame_number = None
181
+ if frame_number is None :
182
+ try :
183
+ frame_number = cvat_task_anno .match_frame (item .id )
184
+ except Exception :
185
+ pass
186
+ if frame_number is None and item .has_image :
187
+ try :
188
+ frame_number = cvat_task_anno .match_frame (item .image .filename )
189
+ except Exception :
190
+ pass
191
+ if frame_number is None :
192
+ try :
193
+ frame_number = int (item .id )
194
+ except Exception :
195
+ pass
196
+ if not frame_number in cvat_task_anno .frame_info :
197
+ raise Exception ("Could not match item id: '%s' with any task frame" %
198
+ item .id )
199
+ return frame_number
200
+
201
+ def import_dm_annotations (dm_dataset , cvat_task_anno ):
202
+ shapes = {
203
+ datumaro .AnnotationType .bbox : ShapeType .RECTANGLE ,
204
+ datumaro .AnnotationType .polygon : ShapeType .POLYGON ,
205
+ datumaro .AnnotationType .polyline : ShapeType .POLYLINE ,
206
+ datumaro .AnnotationType .points : ShapeType .POINTS ,
207
+ }
208
+
209
+ label_cat = dm_dataset .categories ()[datumaro .AnnotationType .label ]
210
+
211
+ for item in dm_dataset :
212
+ frame_number = match_frame (item , cvat_task_anno )
213
+
214
+ for ann in item .annotations :
215
+ if ann .type in shapes :
216
+ cvat_task_anno .add_shape (cvat_task_anno .LabeledShape (
217
+ type = shapes [ann .type ],
218
+ frame = frame_number ,
219
+ label = label_cat .items [ann .label ].name ,
220
+ points = ann .points ,
221
+ occluded = False ,
222
+ attributes = [],
223
+ ))
0 commit comments