plate augment

Text-Image-Augmentation-python-master/demo1.py

@@ -0,0 +1,72 @@
+# -*- coding:utf-8 -*-
+# Author: RubanSeven
+
+import cv2
+import imageio
+import numpy as np
+import os
+from augment import distort, stretch, perspective
+import argparse
+
+def allFileList(rootfile,allFile):
+    folder =os.listdir(rootfile)
+    for temp in folder:
+        fileName = os.path.join(rootfile,temp)
+        if os.path.isfile(fileName):
+            allFile.append(fileName)
+        else:
+            allFileList(fileName,allFile)
+
+def create_gif(image_list, gif_name, duration=0.1):
+    frames = []
+    for image in image_list:
+        frames.append(image)
+    imageio.mimsave(gif_name, frames, 'GIF', duration=duration)
+    return
+
+
+if __name__ == '__main__':
+    parser =argparse.ArgumentParser()
+    parser.add_argument('--src_path', type=str, default='/mnt/Gu/trainData/weixian', help='model.pt path(s)')  
+    parser.add_argument('--dst_path', type=str, default='/mnt/Gu/trainData/plate/final/weixian_aug', help='source') 
+    opt = parser.parse_args()
+    rootFile = opt.src_path
+    saveFile = opt.dst_path
+    fileList = []
+    allFileList(rootFile, fileList)
+    picOunt=0
+    for temp in fileList:
+        print(picOunt,temp)
+        picOunt+=1
+        im = cv2.imdecode(np.fromfile(temp,dtype=np.uint8),-1)
+        _,_,c = im.shape
+        if c!=3:
+            continue
+    # im = cv2.resize(im, (200, 64))
+    #     cv2.imshow("im_CV", im)
+        distort_img_list = list()
+        stretch_img_list = list()
+        perspective_img_list = list()
+        for i in range(1):
+            try:
+                distort_img = distort(im, 8)
+                distort_img_list.append(distort_img)
+                # cv2.imshow("distort_img", distort_img)
+
+                stretch_img = stretch(distort_img, 8)
+                # cv2.imshow("stretch_img", stretch_img)
+                stretch_img_list.append(stretch_img)
+
+                name = temp.split(os.sep)[-1].split(".")[0]
+                name = name+str(picOunt)+".jpg"
+                newPath = os.path.join(saveFile,name)
+                perspective_img = perspective(stretch_img)
+                # cv2.imshow("perspective_img", perspective_img)
+                perspective_img_list.append(perspective_img)
+                # cv2.waitKey(1)
+                cv2.imencode('.jpg', perspective_img)[1].tofile(newPath)
+            except:
+                print("transForm error")
+    # create_gif(distort_img_list, r'imgs/distort.gif')
+    # create_gif(stretch_img_list, r'imgs/stretch.gif')
+    # create_gif(perspective_img_list, r'imgs/perspective.gif')