Spaces:
Runtime error
Runtime error
| import imgaug as ia | |
| from imgaug import augmenters as iaa | |
| import numpy as np | |
| import cv2 | |
| from pascal_voc_writer import Writer | |
| import xml.etree.ElementTree as ET | |
| import glob | |
| from util import sequence_simple, sequence_resize | |
| from util import annotation as an | |
| import shutil | |
| import sys | |
| INPUT_DIR = './images' | |
| OUTPUT_DIR = './augmented' | |
| AUGMENT_SIZE = 10 | |
| def main(): | |
| global INPUT_DIR | |
| for classname in glob.glob(INPUT_DIR + "/"): | |
| print(classname) | |
| INPUT_DIR = classname | |
| print(OUTPUT_DIR) | |
| for file in glob.glob('%s/*.xml' % INPUT_DIR): | |
| print('Augmenting %s ...' % file) | |
| annotation = an.parse_xml(file) | |
| augment(annotation) | |
| for file in glob.glob('%s/*.xml' % OUTPUT_DIR): | |
| an.inspect(file) | |
| def augment(annotation): | |
| seq = sequence_resize.get() | |
| for i in range(AUGMENT_SIZE): | |
| filename = annotation['filename'] | |
| sp = filename.split('.') | |
| outfile = '%s/%s-%02d.%s' % (OUTPUT_DIR, sp[0], i, sp[-1]) | |
| print("++INFO++") | |
| print(outfile) | |
| print("++INFO++") | |
| seq_det = seq.to_deterministic() | |
| print("===") | |
| print(INPUT_DIR) | |
| print(annotation['filename']) | |
| print("===") | |
| image = cv2.imread('%s/%s' % (INPUT_DIR, annotation['filename'])) | |
| _bbs = [] | |
| for obj in annotation['objects']: | |
| bb = ia.BoundingBox(x1=int(obj['xmin']), | |
| y1=int(obj['ymin']), | |
| x2=int(obj['xmax']), | |
| y2=int(obj['ymax']), | |
| label=obj['name']) | |
| _bbs.append(bb) | |
| bbs = ia.BoundingBoxesOnImage(_bbs, shape=image.shape) | |
| image_aug = seq_det.augment_images([image])[0] | |
| bbs_aug = seq_det.augment_bounding_boxes( | |
| [bbs])[0].remove_out_of_image().cut_out_of_image() | |
| writer = Writer(outfile, | |
| annotation['size']['width'], | |
| annotation['size']['height']) | |
| for bb in bbs_aug.bounding_boxes: | |
| if int((bb.x2-bb.x1)*(bb.y2-bb.y1)) == 0: | |
| print("augmentet boundingbox has non existing area. Skipping") | |
| continue | |
| writer.addObject(bb.label, | |
| int(bb.x1), | |
| int(bb.y1), | |
| int(bb.x2), | |
| int(bb.y2)) | |
| print(outfile) | |
| cv2.imwrite(outfile, image_aug) | |
| writer.save('%s.xml' % outfile[:-4]) | |
| # writer.save('%s.xml' % outfile.split('.')[0]) | |
| if __name__ == '__main__': | |
| main() | |