Skip to content

Commit

Permalink
Reimplement logic of getting labels in coco converter (#319)
Browse files Browse the repository at this point in the history 
* Reimplement logic of getting labels in coco converter

* Fixes in converter
  • Loading branch information
@DmitriySidnev @nmanovic
DmitriySidnev authored and nmanovic committed Apr 1, 2019
1 parent e96d4fd commit 21de4cf
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 48 deletions.
4 changes: 2 additions & 2 deletions utils/coco/converter.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,13 @@ $ cat ../requirements.txt requirements.txt | xargs -n 1 -L 1 pip install
Run the script inside the virtual environment.

```bash
python converter.py --cvat-xml </path/to/cvat/annotation.xml> --output </path/to/output/coco/annotation.json> --image-dir </path/to/directory/with/images> --draw </path/to/save/directory> --draw_labels --use_background_label
python converter.py --cvat-xml </path/to/cvat/annotation.xml> --output </path/to/output/coco/annotation.json> --image-dir </path/to/directory/with/images> --labels </path/to/file/with/labels.txt> --draw </path/to/save/directory> --draw_labels --use_background_label
```

Please run `python converter.py --help` for more details.

#### Labels
The script parses input annotation and find field `labels` to find which labels are presented. If were not found any labels in annotation, script try to parse text file `labels.txt` in the same directory as annotation. This file should include labels in one string separated by spaces or one label per string and also their combinations. For example:
If '--labels' argument is used, the script gets names of labels from a file. If file with labels is not defined, the script parses input annotation and find field `labels` to find which labels are presented. File with labels should include labels in one string separated by spaces or one label per string and also their combinations. For example:
```
label1 label2
label3
Expand Down
136 changes: 90 additions & 46 deletions utils/coco/converter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,17 @@ def parse_args():
help='input file with CVAT annotation in *.xml format'
)
parser.add_argument(
'--output', required=True,
help='output annotation file'
'--output', default=None,
help='output annotation file. If not defined, the output file name will be created from input file name'
)
parser.add_argument(
'--image-dir', required=True,
help='directory with images from annotation'
)
parser.add_argument(
'--labels', default=None,
help='path to file with labels'
)
parser.add_argument(
'--draw', default=None,
help='directory to save images with its segments. By default is disabled'
Expand Down Expand Up @@ -89,6 +93,7 @@ def mask_to_polygon(mask, tolerance=1.0, area_threshold=1):
polygons.append(reshaped_contour)
return polygons


def draw_polygons(polygons, img_name, input_dir, output_dir, draw_labels):
"""Draw on image contours of its objects and save
Args:
Expand Down Expand Up @@ -123,6 +128,7 @@ def draw_polygons(polygons, img_name, input_dir, output_dir, draw_labels):
cv2.putText(img, label, (x, y), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, red, 1)
cv2.imwrite(output_file, img)


def fix_segments_intersections(polygons, height, width, img_name, use_background_label,
threshold=0.0, ratio_tolerance=0.001, area_threshold=1):
"""Find all intersected regions and crop contour for back object by objects which
Expand Down Expand Up @@ -206,6 +212,7 @@ def fix_segments_intersections(polygons, height, width, img_name, use_background

return output_polygons


def polygon_area_and_bbox(polygon, height, width):
"""Calculate area of object's polygon and bounding box around it
Args:
Expand All @@ -222,6 +229,7 @@ def polygon_area_and_bbox(polygon, height, width):
max(bbox[:, 1] + bbox[:, 3]) - min(bbox[:, 1])]
return area, bbox


def insert_license_data(result_annotation):
"""Fill license fields in annotation by blank data
Args:
Expand All @@ -233,6 +241,7 @@ def insert_license_data(result_annotation):
'url': ''
})


def insert_info_data(xml_root, result_annotation):
"""Fill available information of annotation
Args:
Expand Down Expand Up @@ -266,56 +275,84 @@ def insert_info_data(xml_root, result_annotation):
}
log.info('Found the next information data: {}'.format(result_annotation['info']))

def insert_categories_data(xml_root, use_background_label, result_annotation, xml_dir):

def insert_categories_data(xml_root, use_background_label, result_annotation, labels_file=None):
"""Get labels from input annotation and fill categories field in output annotation
Args:
xml_root: root for xml parser
use_background_label: key to enable using label background
result_annotation: output annotation in COCO representation
xml_dir: directory with input annotation
labels_file: path to file with labels names.
If not defined, parse annotation to get labels names
"""
log.info('Reading labels...')
def get_categories(names, bg_found, use_background_label, sort=False):
bg_used = False
category_map = {}
categories = []
# Sort labels by its names to make the same order of ids for different annotations
if sort:
names.sort()
# Always use id = 0 for background
if bg_found and use_background_label:
category_map['background'] = 0
bg_used = True
cat_id = 1
# Define id for all labels beginning from 1
for name in names:
if name == 'background':
continue
category_map[name] = cat_id
categories.append({'id': cat_id, 'name': name, 'supercategory': ''})
cat_id += 1
return category_map, categories, bg_used

categories = []
category_map = {}
label_names = []
bg_found = False
id = 0
for label in xml_root.iter('label'):
for name in label.findall("./name"):
if not use_background_label and name.text == 'background':
bg_found = True
continue
category_map[name.text] = id
categories.append({'id': id, 'name': name.text, 'supercategory': ''})
id += 1
bg_used = False

if labels_file is None:
log.info('Reading labels from annotation...')
for label in xml_root.iter('label'):
for name in label.findall("./name"):
if name.text == 'background':
bg_found = True
continue
label_names.append(name.text)
if len(label_names) == 0:
log.info('Labels in annotation were not found. Please use \'--labels\' argument to define file with labels.')
else:
category_map, categories, bg_used = get_categories(label_names, bg_found, use_background_label, sort=True)
else:
log.info('Parsing labels from file <{}>...'.format(labels_file))
with open(labels_file, 'r') as file:
string = ' '
while string != '' and string != '\n':
string = file.readline()
labels = string.split(' ')
for label in labels:
if label == '\n' or label == '':
continue
label = label.replace('\n', '')
if label == 'background':
bg_found = True
continue
label_names.append(label)
category_map, categories, bg_used = get_categories(label_names, bg_found, use_background_label)

if len(categories) == 0:
log.info('Labels in annotation were not found. Trying to find file <labels.txt> in <{}>'.format(xml_dir))
if osp.isfile(osp.join(xml_dir, 'labels.txt')):
labels_file = osp.join(xml_dir, 'labels.txt')
log.info('File <labels.txt> was found in <{}>. Reading...'.format(xml_dir))
with open(labels_file, 'r') as file:
string = ' '
id = 0
while string != '' and string != '\n':
string = file.readline()
labels = string.split(' ')
for l in labels:
if l == '\n':
continue
if not use_background_label and l == 'background':
bg_found = True
continue
category_map[l] = id
categories.append({'id': id, 'name': l, 'supercategory': ''})
id += 1
raise ValueError('Categories list is empty. Something wrong.')

result_annotation['categories'] = categories
log.info('Found the next labels: {}'.format(category_map))
if bg_found:
if bg_found and not bg_used:
log.warning('Label <background> was found but not used. '
'To enable it should use command line argument [--use_background_label]')
return category_map

def insert_image_data(image, path_to_images, result_annotation):

def insert_image_data(image, result_annotation):
"""Get data from input annotation for image and fill fields for this image in output annotation
Args:
image: dictionary with data for image from original annotation
Expand All @@ -329,11 +366,11 @@ def insert_image_data(image, path_to_images, result_annotation):
new_img['license'] = 0
new_img['id'] = image['id']
new_img['file_name'] = osp.basename(image['name'])
pic = cv2.imread(osp.join(path_to_images, new_img['file_name']))
new_img['height'] = pic.shape[0]
new_img['width'] = pic.shape[1]
new_img['height'] = int(image['height'])
new_img['width'] = int(image['width'])
result_annotation['images'].append(new_img)


def insert_annotation_data(image, category_map, segm_id, object, img_dims, result_annotation):
"""Get data from input annotation for object and fill fields for this object in output annotation
Args:
Expand All @@ -359,10 +396,14 @@ def insert_annotation_data(image, category_map, segm_id, object, img_dims, resul
def main():
args = parse_args()
xml_file_name = args.cvat_xml
output_file_name = args.output
if args.output is not None:
output_file_name = args.output
else:
output_file_name = args.cvat_xml.split('.xml')[0] + '.json'
log.info('Output file name set to: {}'.format(output_file_name))
root = etree.parse(xml_file_name).getroot()

if args.draw != None:
if args.draw is not None:
log.info('Draw key was enabled. Images will be saved in directory <{}>'.format(args.draw))

result_annotation = {
Expand All @@ -375,7 +416,7 @@ def main():

insert_license_data(result_annotation)
insert_info_data(root, result_annotation)
category_map = insert_categories_data(root, args.use_background_label, result_annotation, osp.dirname(xml_file_name))
category_map = insert_categories_data(root, args.use_background_label, result_annotation, labels_file=args.labels)

if len(category_map) == 0:
sys.exit('Labels were not found. Be sure that annotation <{}> includes field <labels> or '
Expand All @@ -388,6 +429,9 @@ def main():
image = {}
for key, value in img.items():
image[key] = value
img_name = osp.join(args.image_dir, osp.basename(image['name']))
if not osp.isfile(img_name):
log.warning('Image <{}> is not available'.format(img_name))
image['polygon'] = []
z_order_on_counter = 0
polygon_counter = 0
Expand All @@ -407,7 +451,7 @@ def main():

# Create new image
image['id'] = int(image['id'])
insert_image_data(image, args.image_dir, result_annotation)
insert_image_data(image, result_annotation)
height = result_annotation['images'][-1]['height']
width = result_annotation['images'][-1]['width']
image['polygon'] = fix_segments_intersections(image['polygon'], height, width,
Expand All @@ -420,7 +464,7 @@ def main():
segm_id += 1

# Draw contours of objects on image
if args.draw != None:
if args.draw is not None:
draw_polygons(image['polygon'], image['name'], args.image_dir, args.draw, args.draw_labels)

log.info('Processed images: {}'.format(len(result_annotation['images'])))
Expand All @@ -438,12 +482,12 @@ def main():
# Try to load created annotation via cocoapi
try:
log.info('Trying to load annotation <{}> via cocoapi...'.format(output_file_name))
anno = coco_loader.COCO(output_file_name)
coco_loader.COCO(output_file_name)
except:
raise
else:
log.info('Annotation in COCO representation <{}> created from <{}> successfully!'
.format(output_file_name, xml_file_name))
log.info('Conversion <{}> --> <{}> has finished successfully!'.format(xml_file_name, output_file_name))


if __name__ == "__main__":
main()

0 comments on commit 21de4cf

Please sign in to comment.