Files
crnn_plate_recognition/Text-Image-Augmentation-python-master/demo1.py
2022-11-15 21:05:11 +08:00

73 lines
2.5 KiB
Python

# -*- 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')