AugmentImagesAndXML/AugmentImageImageNetPolicy.py at master · deepdigger/AugmentImagesAndXML · GitHub

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
from autoaugment import ImageNetPolicy
from PIL import Image
import glob
import xml.etree.ElementTree as ET
import StaticMethods as sm

sm.createFolder("AutoAugmentImages")
sm.createFolder("AutoAugmentXML")
policy = ImageNetPolicy()

for XMLPath in glob.glob("XML_Files_Input\\*"):
    fileName = XMLPath.split("\\")[1].split(".")[0]
    SplitAtDotXMLPath = XMLPath.split(".")
    tree = ET.parse(XMLPath)
    tree.write("AutoAugmentXML\\" + fileName + ".xml")
    root = tree.getroot()
    try:
        img = Image.open("Pictures\\" + fileName + ".jpg")
        img.save("AutoAugmentImages\\" + fileName + ".jpg")


        print("Autoaugment: %s" % root.find("filename").text)
        for i in range(15):
            newName = fileName + "AutoAugment" + str(i)
            root.find("filename").text = newName + ".jpg"
            tree.write("AutoAugmentXML\\" + newName + ".xml")
            augmentedImage = policy(img, i)
            augmentedImage.save("AutoAugmentImages\\" + newName + ".jpg")

    except FileNotFoundError:
        print("Image: %s not found" % fileName)