Spaces:
Runtime error
Runtime error
File size: 2,659 Bytes
68896f3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
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()
|