if __name__ == '__main__':
import os
import cv2
from util.augmentation import Augmentation
from util.misc import regularize_sin_cos
from nmslib import lanms
from util.pbox import bbox_transfor_inv, minConnectPath
from util import canvas as cav
import time
means = (0.485, 0.456, 0.406)
stds = (0.229, 0.224, 0.225)
transform = Augmentation(size=640, mean=means, std=stds)
trainset = TD500Text(data_root='../data/TD500',
is_training=False,
transform=transform)
# img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, meta = trainset[944]
for idx in range(0, len(trainset)):
t0 = time.time()
img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi = trainset[
idx]
img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi \
= map(lambda x: x.cpu().numpy(), (img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi))
img = img.transpose(1, 2, 0)
img = ((img * stds + means) * 255).astype(np.uint8)
def main():
global lr
if cfg.exp_name == 'Totaltext':
trainset = TotalText(data_root='data/total-text-mat',
ignore_list=None,
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
# valset = TotalText(
# data_root='data/total-text-mat',
# ignore_list=None,
# is_training=False,
# transform=BaseTransform(size=cfg.input_size, mean=cfg.means, std=cfg.stds)
# )
valset = None
elif cfg.exp_name == 'Synthtext':
trainset = SynthText(data_root='data/SynthText',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
elif cfg.exp_name == 'Ctw1500':
trainset = Ctw1500Text(data_root='data/ctw1500',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
elif cfg.exp_name == 'Icdar2015':
trainset = Icdar15Text(data_root='data/Icdar2015',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
elif cfg.exp_name == 'MLT2017':
trainset = Mlt2017Text(data_root='data/MLT2017',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
elif cfg.exp_name == 'TD500':
trainset = TD500Text(data_root='data/TD500',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
else:
print("dataset name is not correct")
train_loader = data.DataLoader(trainset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
pin_memory=True)
log_dir = os.path.join(
cfg.log_dir,
datetime.now().strftime('%b%d_%H-%M-%S_') + cfg.exp_name)
logger = LogSummary(log_dir)
# Model
model = TextNet(backbone=cfg.net, is_training=True)
if cfg.mgpu:
model = nn.DataParallel(model)
model = model.to(cfg.device)
if cfg.cuda:
cudnn.benchmark = True
if cfg.resume:
load_model(model, cfg.resume)
criterion = TextLoss()
lr = cfg.lr
moment = cfg.momentum
if cfg.optim == "Adam" or cfg.exp_name == 'Synthtext':
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
else:
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=moment)
if cfg.exp_name == 'Synthtext':
scheduler = FixLR(optimizer)
else:
scheduler = lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.90)
print('Start training TextGraph.')
for epoch in range(cfg.start_epoch, cfg.start_epoch + cfg.max_epoch + 1):
scheduler.step()
train(model, train_loader, criterion, scheduler, optimizer, epoch,
logger)
print('End.')
if torch.cuda.is_available():
torch.cuda.empty_cache()
def main():
global lr
trainset = None
valset = None
if cfg.dataset_name == 'total_text':
trainset = TotalText(data_root=cfg.dataset_root,
ignore_list=None,
is_training=True,
transform=Augmentation(size=cfg.img_size,
mean=cfg.means,
std=cfg.stds))
valset = TotalText(data_root=cfg.dataset_root,
ignore_list=None,
is_training=False,
transform=BaseTransform(size=cfg.img_size,
mean=cfg.means,
std=cfg.stds))
elif cfg.dataset_name == 'coco_text':
trainset = COCO_Text(data_root=cfg.dataset_root,
is_training=True,
transform=Augmentation(size=cfg.img_size,
mean=cfg.means,
std=cfg.stds))
valset = None
elif cfg.dataset_name == 'synth_text':
trainset = SynthText(
data_root='/home/andrew/Documents/Dataset/SynthText/SynthText',
is_training=True,
transform=Augmentation(size=cfg.img_size,
mean=cfg.means,
std=cfg.stds))
valset = None
else:
pass
train_loader = data.DataLoader(trainset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.training_num_workers,
pin_memory=True,
timeout=10)
if valset:
val_loader = data.DataLoader(valset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.training_num_workers)
else:
val_loader = None
model = UHT_Net(pretrained=True)
if cfg.multi_gpu:
model = nn.DataParallel(model)
model = model.to(cfg.device)
if cfg.cuda:
cudnn.benchmark = True
if cfg.resume != "":
load_model(model, cfg.resume)
criterion = UHT_Loss()
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=len(train_loader) *
cfg.decay_epoch,
gamma=cfg.decay_rate)
print('Start training Model.')
for epoch in range(cfg.start_epoch, cfg.end_epoch):
train(model, train_loader, criterion, scheduler, optimizer, epoch)
if valset:
validation(model, val_loader, criterion)
print('End.')
def main():
global lr
if cfg.dataset == 'total-text':
trainset = TotalText(data_root='data/total-text',
ignore_list=None,
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = TotalText(data_root='data/total-text',
ignore_list=None,
is_training=False,
transform=BaseTransform(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
elif cfg.dataset == 'synth-text':
trainset = SynthText(data_root='data/SynthText',
is_training=True,
transform=Augmentation(size=cfg.input_size,
mean=cfg.means,
std=cfg.stds))
valset = None
else:
pass
train_loader = data.DataLoader(trainset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers)
if valset:
val_loader = data.DataLoader(valset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers)
else:
valset = None
log_dir = os.path.join(
cfg.log_dir,
datetime.now().strftime('%b%d_%H-%M-%S_') + cfg.exp_name)
logger = LogSummary(log_dir)
# Model
model = TextNet(is_training=True, backbone=cfg.net)
if cfg.mgpu:
model = nn.DataParallel(model)
model = model.to(cfg.device)
if cfg.cuda:
cudnn.benchmark = True
if cfg.resume:
load_model(model, cfg.resume)
criterion = TextLoss()
lr = cfg.lr
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
if cfg.dataset == 'synth-text':
scheduler = FixLR(optimizer)
else:
scheduler = lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.1)
print('Start training TextSnake.')
for epoch in range(cfg.start_epoch, cfg.max_epoch):
train(model, train_loader, criterion, scheduler, optimizer, epoch,
logger)
if valset:
validation(model, val_loader, criterion, epoch, logger)
print('End.')