wangchunlin
/
cvat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Reimplement logic of getting labels in coco converter (#319)

Browse Source 
* Reimplement logic of getting labels in coco converter

* Fixes in converter
main
DmitriySidnev 7 years ago committed by Nikita Manovich
parent e96d4fdab4
commit 21de4cf7b3
2 changed files with 92 additions and 48 deletions
Show Stats Download Patch File Download Diff File

4
utils/coco/converter.md
Unescape Escape View File

@ -24,13 +24,13 @@ $ cat ../requirements.txt requirements.txt | xargs -n 1 -L 1 pip install
Run the script inside the virtual environment. Run the script inside the virtual environment.
```bash ```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. Please run `python converter.py --help` for more details.
#### Labels #### 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 label1 label2
label3 label3

136
utils/coco/converter.py
Unescape Escape View File

@ -30,13 +30,17 @@ def parse_args():
help='input file with CVAT annotation in *.xml format' help='input file with CVAT annotation in *.xml format'
) )
parser.add_argument( parser.add_argument(
'--output', required=True, '--output', default=None,
help='output annotation file' help='output annotation file. If not defined, the output file name will be created from input file name'
) )
parser.add_argument( parser.add_argument(
'--image-dir', required=True, '--image-dir', required=True,
help='directory with images from annotation' help='directory with images from annotation'
) )
parser.add_argument(
'--labels', default=None,
help='path to file with labels'
)
parser.add_argument( parser.add_argument(
'--draw', default=None, '--draw', default=None,
help='directory to save images with its segments. By default is disabled' help='directory to save images with its segments. By default is disabled'
@ -89,6 +93,7 @@ def mask_to_polygon(mask, tolerance=1.0, area_threshold=1):
polygons.append(reshaped_contour) polygons.append(reshaped_contour)
return polygons return polygons
def draw_polygons(polygons, img_name, input_dir, output_dir, draw_labels): def draw_polygons(polygons, img_name, input_dir, output_dir, draw_labels):
"""Draw on image contours of its objects and save """Draw on image contours of its objects and save
Args: Args:
@ -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.putText(img, label, (x, y), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, red, 1)
cv2.imwrite(output_file, img) cv2.imwrite(output_file, img)
def fix_segments_intersections(polygons, height, width, img_name, use_background_label, def fix_segments_intersections(polygons, height, width, img_name, use_background_label,
threshold=0.0, ratio_tolerance=0.001, area_threshold=1): threshold=0.0, ratio_tolerance=0.001, area_threshold=1):
"""Find all intersected regions and crop contour for back object by objects which """Find all intersected regions and crop contour for back object by objects which
@ -206,6 +212,7 @@ def fix_segments_intersections(polygons, height, width, img_name, use_background
return output_polygons return output_polygons
def polygon_area_and_bbox(polygon, height, width): def polygon_area_and_bbox(polygon, height, width):
"""Calculate area of object's polygon and bounding box around it """Calculate area of object's polygon and bounding box around it
Args: Args:
@ -222,6 +229,7 @@ def polygon_area_and_bbox(polygon, height, width):
max(bbox[:, 1] + bbox[:, 3]) - min(bbox[:, 1])] max(bbox[:, 1] + bbox[:, 3]) - min(bbox[:, 1])]
return area, bbox return area, bbox
def insert_license_data(result_annotation): def insert_license_data(result_annotation):
"""Fill license fields in annotation by blank data """Fill license fields in annotation by blank data
Args: Args:
@ -233,6 +241,7 @@ def insert_license_data(result_annotation):
'url': '' 'url': ''
}) })
def insert_info_data(xml_root, result_annotation): def insert_info_data(xml_root, result_annotation):
"""Fill available information of annotation """Fill available information of annotation
Args: Args:
@ -266,56 +275,84 @@ def insert_info_data(xml_root, result_annotation):
} }
log.info('Found the next information data: {}'.format(result_annotation['info'])) 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 """Get labels from input annotation and fill categories field in output annotation
Args: Args:
xml_root: root for xml parser xml_root: root for xml parser
use_background_label: key to enable using label background use_background_label: key to enable using label background
result_annotation: output annotation in COCO representation 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 = [] categories = []
category_map = {} category_map = {}
label_names = []
bg_found = False bg_found = False
id = 0 bg_used = False
for label in xml_root.iter('label'):
for name in label.findall("./name"): if labels_file is None:
if not use_background_label and name.text == 'background': log.info('Reading labels from annotation...')
bg_found = True for label in xml_root.iter('label'):
continue for name in label.findall("./name"):
category_map[name.text] = id if name.text == 'background':
categories.append({'id': id, 'name': name.text, 'supercategory': ''}) bg_found = True
id += 1 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: if len(categories) == 0:
log.info('Labels in annotation were not found. Trying to find file <labels.txt> in <{}>'.format(xml_dir)) raise ValueError('Categories list is empty. Something wrong.')
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
result_annotation['categories'] = categories result_annotation['categories'] = categories
log.info('Found the next labels: {}'.format(category_map)) 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. ' log.warning('Label <background> was found but not used. '
'To enable it should use command line argument [--use_background_label]') 'To enable it should use command line argument [--use_background_label]')
return category_map 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 """Get data from input annotation for image and fill fields for this image in output annotation
Args: Args:
image: dictionary with data for image from original annotation image: dictionary with data for image from original annotation
@ -329,11 +366,11 @@ def insert_image_data(image, path_to_images, result_annotation):
new_img['license'] = 0 new_img['license'] = 0
new_img['id'] = image['id'] new_img['id'] = image['id']
new_img['file_name'] = osp.basename(image['name']) new_img['file_name'] = osp.basename(image['name'])
pic = cv2.imread(osp.join(path_to_images, new_img['file_name'])) new_img['height'] = int(image['height'])
new_img['height'] = pic.shape[0] new_img['width'] = int(image['width'])
new_img['width'] = pic.shape[1]
result_annotation['images'].append(new_img) result_annotation['images'].append(new_img)
def insert_annotation_data(image, category_map, segm_id, object, img_dims, result_annotation): 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 """Get data from input annotation for object and fill fields for this object in output annotation
Args: Args:
@ -359,10 +396,14 @@ def insert_annotation_data(image, category_map, segm_id, object, img_dims, resul
def main(): def main():
args = parse_args() args = parse_args()
xml_file_name = args.cvat_xml 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() 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)) log.info('Draw key was enabled. Images will be saved in directory <{}>'.format(args.draw))
result_annotation = { result_annotation = {
@ -375,7 +416,7 @@ def main():
insert_license_data(result_annotation) insert_license_data(result_annotation)
insert_info_data(root, 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: if len(category_map) == 0:
sys.exit('Labels were not found. Be sure that annotation <{}> includes field <labels> or ' sys.exit('Labels were not found. Be sure that annotation <{}> includes field <labels> or '
@ -388,6 +429,9 @@ def main():
image = {} image = {}
for key, value in img.items(): for key, value in img.items():
image[key] = value 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'] = [] image['polygon'] = []
z_order_on_counter = 0 z_order_on_counter = 0
polygon_counter = 0 polygon_counter = 0
@ -407,7 +451,7 @@ def main():
# Create new image # Create new image
image['id'] = int(image['id']) 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'] height = result_annotation['images'][-1]['height']
width = result_annotation['images'][-1]['width'] width = result_annotation['images'][-1]['width']
image['polygon'] = fix_segments_intersections(image['polygon'], height, width, image['polygon'] = fix_segments_intersections(image['polygon'], height, width,
@ -420,7 +464,7 @@ def main():
segm_id += 1 segm_id += 1
# Draw contours of objects on image # 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) draw_polygons(image['polygon'], image['name'], args.image_dir, args.draw, args.draw_labels)
log.info('Processed images: {}'.format(len(result_annotation['images']))) log.info('Processed images: {}'.format(len(result_annotation['images'])))
@ -438,12 +482,12 @@ def main():
# Try to load created annotation via cocoapi # Try to load created annotation via cocoapi
try: try:
log.info('Trying to load annotation <{}> via cocoapi...'.format(output_file_name)) 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: except:
raise raise
else: else:
log.info('Annotation in COCO representation <{}> created from <{}> successfully!' log.info('Conversion <{}> --> <{}> has finished successfully!'.format(xml_file_name, output_file_name))
.format(output_file_name, xml_file_name))
if __name__ == "__main__": if __name__ == "__main__":
main() main()

Write Preview
Loading…
Cancel
Save