@ -14,7 +14,7 @@ from datumaro.components.converter import Converter
from datumaro . components . extractor import (
DEFAULT_SUBSET_NAME , AnnotationType , PointsObject , BboxObject
)
from datumaro . components . formats . ms_coco import Coco AnnotationType , CocoPath
from datumaro . components . formats . ms_coco import Coco Task , CocoPath
from datumaro . util import find
from datumaro . util . image import save_image
import datumaro . util . mask_tools as mask_tools
@ -29,8 +29,9 @@ def _cast(value, type_conv, default=None):
return default
class _TaskConverter :
def __init__ ( self ):
def __init__ ( self , context ):
self . _min_ann_id = 1
self . _context = context
data = {
' licenses ' : [ ] ,
@ -191,6 +192,13 @@ class _InstancesConverter(_TaskConverter):
rle = mask_utils . merge ( rles )
area = mask_utils . area ( rle )
if self . _context . _merge_polygons :
binary_mask = mask_utils . decode ( rle ) . astype ( np . bool )
binary_mask = np . asfortranarray ( binary_mask , dtype = np . uint8 )
segmentation = mask_tools . convert_mask_to_rle ( binary_mask )
is_crowd = True
bbox = [ int ( i ) for i in mask_utils . toBbox ( rle ) ]
if ann . group is not None :
# Mark the group as visited to prevent repeats
for a in annotations [ : ] :
@ -201,6 +209,18 @@ class _InstancesConverter(_TaskConverter):
is_crowd = False
segmentation = [ ann . get_polygon ( ) ]
area = ann . area ( )
if self . _context . _merge_polygons :
h , w , _ = item . image . shape
rles = mask_utils . frPyObjects ( segmentation , h , w )
rle = mask_utils . merge ( rles )
area = mask_utils . area ( rle )
binary_mask = mask_utils . decode ( rle ) . astype ( np . bool )
binary_mask = np . asfortranarray ( binary_mask , dtype = np . uint8 )
segmentation = mask_tools . convert_mask_to_rle ( binary_mask )
is_crowd = True
bbox = [ int ( i ) for i in mask_utils . toBbox ( rle ) ]
if bbox is None :
bbox = ann . get_bbox ( )
@ -340,22 +360,30 @@ class _LabelsConverter(_TaskConverter):
class _Converter :
_TASK_CONVERTER = {
Coco AnnotationType . image_info : _ImageInfoConverter ,
Coco AnnotationType . instances : _InstancesConverter ,
Coco AnnotationType . person_keypoints : _KeypointsConverter ,
Coco AnnotationType . captions : _CaptionsConverter ,
Coco AnnotationType . labels : _LabelsConverter ,
Coco Task . image_info : _ImageInfoConverter ,
Coco Task . instances : _InstancesConverter ,
Coco Task . person_keypoints : _KeypointsConverter ,
Coco Task . captions : _CaptionsConverter ,
Coco Task . labels : _LabelsConverter ,
}
def __init__ ( self , extractor , save_dir , save_images = False , task = None ) :
if not task :
task = list ( self . _TASK_CONVERTER . keys ( ) )
elif task in CocoAnnotationType :
task = [ task ]
self . _task = task
def __init__ ( self , extractor , save_dir ,
tasks = None , save_images = False , merge_polygons = False ) :
assert tasks is None or isinstance ( tasks , ( CocoTask , list ) )
if tasks is None :
tasks = list ( self . _TASK_CONVERTER )
elif isinstance ( tasks , CocoTask ) :
tasks = [ tasks ]
else :
for t in tasks :
assert t in CocoTask
self . _tasks = tasks
self . _extractor = extractor
self . _save_dir = save_dir
self . _save_images = save_images
self . _merge_polygons = merge_polygons
def make_dirs ( self ) :
self . _images_dir = osp . join ( self . _save_dir , CocoPath . IMAGES_DIR )
@ -365,11 +393,13 @@ class _Converter:
os . makedirs ( self . _ann_dir , exist_ok = True )
def make_task_converter ( self , task ) :
return self . _TASK_CONVERTER [ task ] ( )
if task not in self . _TASK_CONVERTER :
raise NotImplementedError ( )
return self . _TASK_CONVERTER [ task ] ( self )
def make_task_converters ( self ) :
return {
task : self . make_task_converter ( task ) for task in self . _task
task : self . make_task_converter ( task ) for task in self . _task s
}
def save_image ( self , item , filename ) :
@ -411,32 +441,56 @@ class _Converter:
' %s _ %s .json ' % ( task . name , subset_name ) ) )
class CocoConverter ( Converter ) :
def __init__ ( self , task = None , save_images = False ) :
def __init__ ( self ,
tasks = None , save_images = False , merge_polygons = False ,
cmdline_args = None ) :
super ( ) . __init__ ( )
self . _task = task
self . _save_images = save_images
self . _options = {
' tasks ' : tasks ,
' save_images ' : save_images ,
' merge_polygons ' : merge_polygons ,
}
if cmdline_args is not None :
self . _options . update ( self . _parse_cmdline ( cmdline_args ) )
@staticmethod
def _split_tasks_string ( s ) :
return [ CocoTask [ i . strip ( ) ] for i in s . split ( ' , ' ) ]
@classmethod
def build_cmdline_parser ( cls , parser = None ) :
import argparse
if not parser :
parser = argparse . ArgumentParser ( )
parser . add_argument ( ' --save-images ' , action = ' store_true ' ,
help = " Save images (default: %(default)s ) " )
parser . add_argument ( ' --merge-polygons ' , action = ' store_true ' ,
help = " Merge instance polygons into a mask (default: %(default)s ) " )
parser . add_argument ( ' --tasks ' , type = cls . _split_tasks_string ,
default = None ,
help = " COCO task filter, comma-separated list of { %s } "
" (default: all) " % ' , ' . join ( [ t . name for t in CocoTask ] ) )
return parser
def __call__ ( self , extractor , save_dir ) :
converter = _Converter ( extractor , save_dir ,
save_images = self . _save_images , task = self . _task )
converter = _Converter ( extractor , save_dir , * * self . _options )
converter . convert ( )
def CocoInstancesConverter ( save_images = False ) :
return CocoConverter ( CocoAnnotationType . instances ,
save_images = save_images )
def CocoInstancesConverter ( * * kwargs ) :
return CocoConverter ( CocoTask . instances , * * kwargs )
def CocoImageInfoConverter ( save_images = False ) :
return CocoConverter ( CocoAnnotationType . image_info ,
save_images = save_images )
def CocoImageInfoConverter ( * * kwargs ) :
return CocoConverter ( CocoTask . image_info , * * kwargs )
def CocoPersonKeypointsConverter ( save_images = False ) :
return CocoConverter ( CocoAnnotationType . person_keypoints ,
save_images = save_images )
def CocoPersonKeypointsConverter ( * * kwargs ) :
return CocoConverter ( CocoTask . person_keypoints , * * kwargs )
def CocoCaptionsConverter ( save_images = False ) :
return CocoConverter ( CocoAnnotationType . captions ,
save_images = save_images )
def CocoCaptionsConverter ( * * kwargs ) :
return CocoConverter ( CocoTask . captions , * * kwargs )
def CocoLabelsConverter ( save_images = False ) :
return CocoConverter ( CocoAnnotationType . labels ,
save_images = save_images )
def CocoLabelsConverter ( * * kwargs ) :
return CocoConverter ( CocoTask . labels , * * kwargs )