def _get_batch(self):
self._batch = self.rec.next()
if not self._batch:
return False
# modify label
label = self._batch.label[0].asnumpy()
for i in range(len(label)):
#if i==0:
#print(label[i])
for j in range(len(label[i])):
if label[i][j] == 78:
label[i][j] = 0
else:
label[i][j] += 1
self._batch.label = [mx.nd.array(label)]
buf = self._batch.data[0].asnumpy()
data = np.zeros((self.batch_size, self.data_shape[0],
self.data_shape[1], self.data_shape[2]))
for i in range(len(buf)):
im = buf[i] #
augment_type = random.randint(0, 2)
ori_im = im.transpose(1, 2, 0).astype(np.uint8)
im = cv2.imread("/opt/data/plate/im.jpg")
# if augment_type == 0:
# im_Distort = Augment.GenerateDistort(im, 12)
# cv2.imwrite("/opt/data/plate/im_Distort.jpg", im_Distort)
# if augment_type == 1:
im_Stretch = Augment.GenerateStretch(im, 12)
cv2.imwrite("/opt/data/plate/im_Stretch.jpg", im_Stretch)
# if augment_type == 2:
im_Perspective = Augment.GeneratePerspective(im)
cv2.imwrite("/opt/data/plate/im_Perspective.jpg", im_Perspective)
data[i] = im.transpose((0, 2, 1))
# cv2.imwrite("/opt/data/plate/im.jpg", im.transpose((1,2,0)))
'''
if i==0:
print(data[i].shape)
#cv2.imshow("e",data[i])
cv2.imwrite("res3.bmp",data[i].transpose((1,2,0)))
cv2.waitKey(0)'''
self._batch.data = [mx.nd.array(data)]
return True
def generate_line(self):
if self.training:
idx = np.random.randint(self.len)
image = self.image[idx]
label = self.label[idx]
else:
idx = self.idx
image = self.image[idx]
label = self.label[idx]
self.idx += 1
if self.idx == self.len:
self.idx -= self.len
h, w = image.shape
imageN = np.ones((self.conH, self.conW)) * 255
beginH = int(abs(self.conH - h) / 2)
beginW = int(abs(self.conW - w) / 2)
if h <= self.conH and w <= self.conW:
imageN[beginH:beginH + h, beginW:beginW + w] = image
elif float(h) / w > float(self.conH) / self.conW:
newW = int(w * self.conH / float(h))
beginW = int(abs(self.conW - newW) / 2)
image = cv2.resize(image, (newW, self.conH))
imageN[:, beginW:beginW + newW] = image
elif float(h) / w <= float(self.conH) / self.conW:
newH = int(h * self.conW / float(w))
beginH = int(abs(self.conH - newH) / 2)
image = cv2.resize(image, (self.conW, newH))
imageN[beginH:beginH + newH] = image
label = self.label[idx]
if self.augment and self.training:
imageN = imageN.astype('uint8')
if torch.rand(1) < 0.3:
imageN = Augment.GenerateDistort(imageN, random.randint(3, 8))
if torch.rand(1) < 0.3:
imageN = Augment.GenerateStretch(imageN, random.randint(3, 8))
if torch.rand(1) < 0.3:
imageN = Augment.GeneratePerspective(imageN)
imageN = imageN.astype('float32')
imageN = (imageN - 127.5) / 127.5
return imageN, label
def word_generate(self):
endW = np.random.randint(50)
label = ''
imageN = np.ones((self.conH, self.conW)) * 255
imageList = []
while True:
idx = np.random.randint(self.len)
image = self.image[idx]
h, w = image.shape
beginH = int(abs(self.conH - h) / 2)
imageList.append(image)
if endW + w > self.conW:
break
if h <= self.conH:
imageN[beginH:beginH + h, endW:endW + w] = image
else:
imageN[:, endW:endW + w] = image[beginH:beginH + self.conH]
endW += np.random.randint(60) + 20 + w
if label == '':
label = self.label[idx]
else:
label = label + '|' + self.label[idx]
label = label
imageN = imageN.astype('uint8')
if self.augment:
if torch.rand(1) < 0.3:
imageN = Augment.GenerateDistort(imageN, random.randint(3, 8))
if torch.rand(1) < 0.3:
imageN = Augment.GenerateStretch(imageN, random.randint(3, 8))
if torch.rand(1) < 0.3:
imageN = Augment.GeneratePerspective(imageN)
imageN = imageN.astype('float32')
imageN = (imageN - 127.5) / 127.5
return imageN, label
import Augment
import cv2
import numpy as np
'''
The code is for OpenCV format.
If your data format is PIL.Image, please convert the format by:
import numpy as np
import cv2
from PIL import Image
img = Image.open("The Path to the image")
img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)
'''
im = cv2.imread("pic/demo.png")
im = cv2.resize(im, (200, 64))
cv2.imshow("im_CV", im)
for i in range(5000):
im_Distort = Augment.GenerateDistort(im, 4)
cv2.imshow("im_Distort", im_Distort)
im_Stretch = Augment.GenerateStretch(im, 4)
cv2.imshow("im_Stretch", im_Stretch)
im_Perspective = Augment.GeneratePerspective(im)
cv2.imshow("im_Perspective", im_Perspective)
cv2.waitKey(1)
def generate_line(self):
# random select in training dataset, and order in testing dataset
if self.is_train:
idx = int(torch.randint(high = self.len,size = (1,)).item())
else:
idx = self.testidx
self.testidx += 1
if self.testidx == self.len:
self.testidx = 0
path = self.path[idx]
image = self.image[idx]
if BINARY:
if KEEP_RATIO:
h,w = image.shape
imageN = np.zeros((self.conH,self.conW))
if self.is_train:
beginH = random.randint(0, int(abs(self.conH-h)))
beginW = random.randint(0, int(abs(self.conW-w)))
else:
beginH = int(abs(self.conH-h)/2)
beginW = int(abs(self.conW-w)/2)
if h <= self.conH and w <= self.conW:
imageN[beginH:beginH+h, beginW:beginW+w] = image
elif float(h) / w > float(self.conH) / self.conW:
newW = int(w * self.conH / float(h))
if self.is_train:
beginW = random.randint(0, int(abs(self.conW-w)))
beginW = int(abs(self.conW-newW)/2)
image = cv2.resize(image, (newW, self.conH))
imageN[:,beginW:beginW+newW] = image
elif float(h) / w <= float(self.conH) / self.conW:
newH = int(h * self.conW / float(w))
if self.is_train:
beginW = random.randint(0, int(abs(self.conW-w)))
beginH = int(abs(self.conH-newH)/2)
image = cv2.resize(image, (self.conW, newH))
imageN[beginH:beginH+newH,:] = image
else:
print("????????????????????????")
else:
image = cv2.resize(image, (self.conW, self.conH))
else:
if KEEP_RATIO:
h,w,_ = image.shape
imageN = np.zeros((self.conH,self.conW,3))
if self.is_train:
beginH = random.randint(0, int(abs(self.conH-h)))
beginW = random.randint(0, int(abs(self.conW-w)))
else:
beginH = int(abs(self.conH-h)/2)
beginW = int(abs(self.conW-w)/2)
if h <= self.conH and w <= self.conW:
imageN[beginH:beginH+h, beginW:beginW+w, :] = image
elif float(h) / w > float(self.conH) / self.conW:
newW = int(w * self.conH / float(h))
if self.is_train:
beginW = random.randint(0, int(abs(self.conW-w)))
beginW = int(abs(self.conW-newW)/2)
image = cv2.resize(image, (newW, self.conH))
imageN[:,beginW:beginW+newW,:] = image
elif float(h) / w <= float(self.conH) / self.conW:
newH = int(h * self.conW / float(w))
if self.is_train:
beginW = random.randint(0, int(abs(self.conW-w)))
beginH = int(abs(self.conH-newH)/2)
image = cv2.resize(image, (self.conW, newH))
imageN[beginH:beginH+newH,:,:] = image
else:
print("????????????????????????")
else:
image = cv2.resize(image, (self.conW, self.conH))
label = self.label[idx]
if AUG:
if self.is_train:
# parts = 4
parts = random.randint(2, 6)
imageN = imageN.astype('uint8')
if torch.rand(1) < AUG_RATIO:
imageN = Augment.GenerateDistort(imageN, parts)
if torch.rand(1) < AUG_RATIO:
imageN = Augment.GenerateStretch(imageN, parts)
if torch.rand(1) < AUG_RATIO:
imageN = Augment.GeneratePerspective(imageN)
if MORE_AUG and self.is_train:
imageN = imageN.astype('uint8')
if np.random.rand() < MORE_AUG_RATIO:
imageN = rot(imageN, r(60) - 30, imageN.shape, 30);
if np.random.rand() < MORE_AUG_RATIO:
imageN = rotRandrom(imageN, 10, (imageN.shape[1], imageN.shape[0]));
if np.random.rand() < MORE_AUG_RATIO:
imageN = AddSmudginess(imageN, self.smu)
if np.random.rand() < MORE_AUG_RATIO and not BINARY:
imageN = tfactor(imageN)
if np.random.rand() < MORE_AUG_RATIO:
imageN = AddGauss(imageN, 1 + r(2));
imageN = cv2.resize(imageN, (self.conW, self.conH))
if torch.rand(1) < 0.001 and self.is_train:
img = cv2.cvtColor(imageN, cv2.COLOR_RGB2BGR)
cv2.imwrite("img_show_train/%s.jpg" % label, img)
return imageN, label, path