Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reimplement logic of getting labels in coco converter #319

Merged
merged 2 commits into from
Apr 1, 2019
Merged

Reimplement logic of getting labels in coco converter #319

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6f55642
Reimplement logic of getting labels in coco converter
Feb 11, 2019
867fbb6
Fixes in converter
Feb 18, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions utils/coco/converter.md
View file
Open in desktop
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
Open in desktop
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 @@ -80,6 +84,7 @@ def mask_to_polygon(mask, tolerance=1.0):
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 @@ -114,6 +119,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):
"""Find all intersected regions and crop contour for back object by objects which
Expand Down Expand Up @@ -197,6 +203,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 @@ -213,6 +220,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 @@ -224,6 +232,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 @@ -257,56 +266,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 @@ -320,11 +357,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 @@ -350,10 +387,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 @@ -366,7 +407,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 @@ -379,6 +420,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 @@ -398,7 +442,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 @@ -410,7 +454,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 @@ -428,12 +472,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()