def __init__(self,
text_file,
image_folder=config.image_folder,
size=config.image_input_size):
with open(text_file, 'r') as f:
self.file_list = f.readlines()
self.length = len(self.file_list)
self.alphabets = Alphabets(config.alphabets)
self.size = size
self.image_folder = image_folder
self.split_magic = " "
def __init__(self,
dataset_dir,
size=config.image_input_size,
transform=default_input_transform,
noise=False):
super(DatasetFromFolder, self).__init__()
self.size = size
self.image_filenames = Path(dataset_dir).rglob("*.png")
self.image_filenames = [item for item in self.image_filenames]
self.alphabets = Alphabets(config.alphabets)
self.transform = transform
self.noise = noise
class DatasetFromTextFile(Dataset):
def __init__(self,
text_file,
image_folder=config.image_folder,
size=config.image_input_size):
with open(text_file, 'r') as f:
self.file_list = f.readlines()
self.length = len(self.file_list)
self.alphabets = Alphabets(config.alphabets)
self.size = size
self.image_folder = image_folder
self.split_magic = "**_~**"
def __getitem__(self, index):
file_name = str(self.file_list[index])
file_name, label = file_name.split(self.split_magic)
label = label.strip()
file_name = os.path.join(self.image_folder, file_name)
image = cv2.imread(file_name, 0)
# image = cv2.resize(image,self.size)
image = input_transform(image)
target = torch.tensor(self.alphabets.decode(label))
return image, target
def __len__(self):
return self.length
class DatasetFromFolder(Dataset):
def __init__(self,
dataset_dir,
size=config.image_input_size,
transform=default_input_transform,
noise=False):
super(DatasetFromFolder, self).__init__()
self.size = size
self.image_filenames = Path(dataset_dir).rglob("*.png")
self.image_filenames = [item for item in self.image_filenames]
self.alphabets = Alphabets(config.alphabets)
self.transform = transform
self.noise = noise
def __getitem__(self, index):
target = str(self.image_filenames[index])
#print(target)
target = target.split('/')[-1]
target = target.split('.')[0]
target = target.split('_')[0]
# target = target.split('/')[-1].split('.')[0].split('_')[0]
image = cv2.imread(str(self.image_filenames[index]))
image = cv2.resize(image, self.size)
if self.noise:
image = salt_and_pepper_noise(image)
image = Image.fromarray(image, mode="RGB")
image = self.transform(image)
#image = self.transform(Image.open(self.image_filenames[index]))
target = torch.tensor(self.alphabets.decode(target))
return image, target
def __len__(self):
return len(self.image_filenames)
class DatasetFromList(Dataset):
def __init__(self, file_list, size=config.image_input_size):
self.file_list = file_list
self.length = len(self.file_list)
self.alphabets = Alphabets(config.alphabets)
self.size = size
def __getitem__(self, index):
target = str(self.file_list[index])
target = target.split('/')[-1].split('_')[0]
image = cv2.imread(str(self.file_list[index]))
image = cv2.resize(image, self.size)
image = input_transform(image)
target = torch.tensor(self.alphabets.decode(target))
return image, target
def __len__(self):
return self.length
class DatasetFromFolder(Dataset):
def __init__(self, dataset_dir, size=config.image_input_size):
super(DatasetFromFolder, self).__init__()
self.size = size
self.image_filenames = Path(dataset_dir).rglob("*.png")
self.image_filenames = [item for item in self.image_filenames]
self.alphabets = Alphabets(config.alphabets)
def __getitem__(self, index):
target = str(self.image_filenames[index])
#target = target.split('/')[-1].split('_')[0]
target = target.split('/')[-1].split('.')[0].split('_')[0]
image = cv2.imread(str(self.image_filenames[index]))
image = cv2.resize(image, self.size)
image = input_transform(image)
#image = self.transform(Image.open(self.image_filenames[index]))
target = torch.tensor(self.alphabets.decode(target))
return image, target
def __len__(self):
return len(self.image_filenames)
# prepare the network
net = Simple_CRNN(len(config.alphabets), input_channels=3, hidden_unit=128)
# prepare the optim
#optimizer = torch.optim.SGD(net.parameters(),args.lr)
#optimizer = torch.optim.SGD(parameters_settings, args.lr)
optimizer = torch.optim.Adam(net.parameters(), config.lr)
#scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=0)
#scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.3)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
milestones=config.step,
gamma=0.1)
# prepare the Alphabets
alphabets = Alphabets(config.alphabets)
# prepare the tensor process
tensor_process = TensorProcess(alphabets)
def get_lr(optimizer):
for param_group in optimizer.param_groups:
return param_group['lr']
def train_epoch(net, epoch):
net = model_to_device(net)
#warmup_scheduler = warmup_controller(optimizer)
for e in range(epoch):
#warmup_scheduler.step(epoch,optimizer)
net.train()
def __init__(self, file_list, size=config.image_input_size):
self.file_list = file_list
self.length = len(self.file_list)
self.alphabets = Alphabets(config.alphabets)
self.size = size
def __init__(self, dataset_dir, size=config.image_input_size):
super(DatasetFromFolder, self).__init__()
self.size = size
self.image_filenames = Path(dataset_dir).rglob("*.png")
self.image_filenames = [item for item in self.image_filenames]
self.alphabets = Alphabets(config.alphabets)