parser = argparse.ArgumentParser(description='DBNet')

parser.add_argument('--config_file',

default='/red_detection/DBNet/DBNet_fzh/config/icdar2015_resnet18_FPN_DBhead_polyLR_word_industry_kd.yaml',

type=str)

parser.add_argument('--local_rank', dest='local_rank', default=0, type=int, help='Use distributed training')

args = parser.parse_args()

training_masks = training_masks.data.cpu().numpy()

pred_text = texts.data.cpu().numpy() * training_masks

pred_text[pred_text <= thred] = 0

pred_text[pred_text > thred] = 1

pred_text = pred_text.astype(np.int32)

gt_text = gt_texts.data.cpu().numpy() * training_masks

gt_text = gt_text.astype(np.int32)

y_true = gt_text.flatten()

y_pred = pred_text.flatten()

precision = precision_score(y_true, y_pred)

recall = recall_score(y_true, y_pred)

f1 = f1_score(y_true, y_pred)

pred_text = preds.astype(np.int32)

gt_text = gt_texts.data.cpu().numpy()

gt_text = gt_text.astype(np.int32)

y_true = gt_text.flatten()

y_pred = pred_text.flatten()

precision = precision_score(y_true, y_pred)

recall = recall_score(y_true, y_pred)

f1 = f1_score(y_true, y_pred)

"""

batch_pred_masks = []

for pred_boxs in boxes_batch:

pred_black = np.zeros((img_h, img_w))

for pred_box in pred_boxs:

pred_box = np.array(pred_box).astype(np.int32)

cv2.fillPoly(pred_black, [pred_box], color=1)

batch_pred_masks.append(pred_black)

"""

if not os.path.exists(output_path):

os.makedirs(output_path)

student_model.eval()

with torch.no_grad():

P, R, F1 = 0, 0, 0

for i, batch in enumerate(validate_loader):

pbar = Progbar(target=len(validate_loader))

# 数据进行转换和丢到gpu

batch_size, _, img_h, img_w = batch['img'].shape

for key, value in batch.items():

if value is not None:

if isinstance(value, torch.Tensor):

batch[key] = value.cuda()

preds = student_model(batch['img'])

batch['shape'] = [(img_h, img_w)] * batch_size

# (batch, nums, 4, 2)

pred_boxes_batch, scores_batch = post_process(batch, preds, is_output_polygon=False)

# print('boxes_batch.shape', boxes_batch)

batch_pred_masks = get_mask(pred_boxes_batch, img_h, img_w)

precision, recall, f1 = get_eval_f1_score(batch_pred_masks, batch['gt'])

P += precision

R += recall

F1 += f1

pbar.update(i + 1, values=[('P', P / (i + 1)), ('R', R / (i + 1)), ('F1', F1 / (i + 1))])

save_model(student_model, optimizer, model_path=os.path.join(output_path,

'model_{}.pth'.format(

str(round(F1 / (i + 1), 3)))),

cycle = np.floor(1 + clr_iterations / (2 * step_size))

x = np.abs(clr_iterations / step_size - 2 * cycle + 1)

lr = base_lr + (max_lr - base_lr) * np.maximum(0, (1 - x)) / (2 ** (cycle - 1))

for param_group in optimizer.param_groups:

param_group['lr'] = lr

post_process, validate_loader, output_path):

for epoch_index in range(epochs):

student_model.train()

pbar = Progbar(target=len(train_loader))

index_train = epoch_index * len(train_loader)

train_loss = 0.0

P, R, F1 = 0, 0, 0

for batch_index, batch in enumerate(train_loader):

batch_index_ = batch_index

batch_index_ += index_train

# lr = optimizer.param_groups[0]['lr']

lr = ajust_learning_tri(optimizer, batch_index_, step_size=len(train_loader) * 8)

# 数据进行转换和丢到gpu

for key, value in batch.items():

if value is not None:

if isinstance(value, torch.Tensor):

batch[key] = value.cuda()

student_preds = student_model(batch['img'])

loss_dict = {}

# KD loss

if teacher_model is not None:

teacher_model.eval()

with torch.no_grad():

# (b,2,h,w)

teacher_outputs = teacher_model(batch['img']) # shrink_maps, threshold_maps, binary_maps

kd_loss_dict = teacher_criterion(student_preds, teacher_outputs, batch)

loss_dict = {**kd_loss_dict}

gt_loss_dict = student_criterion(student_preds, batch)

loss_dict = {**gt_loss_dict, **loss_dict}

# backward

total_losses = sum(loss_ for loss_ in loss_dict.values())

optimizer.zero_grad()

total_losses.backward()

optimizer.step()

precision, recall, f1 = get_f1_score(student_preds[:, 0, :, :], batch['shrink_map'], batch['shrink_mask'],

config['post_processing']['args']['thresh'])

train_loss += total_losses.item()

P += precision

R += recall

F1 += f1

pbar.update(batch_index + 1, values=[('loss', train_loss / (batch_index + 1)),

('P', P / (batch_index + 1)),

('R', R / (batch_index + 1)),

('F1', F1 / (batch_index + 1)),

('epoch:', epoch_index)])

lr_list.append(lr)

if (epoch_index + 1) % 10 == 0:

data_path = config['dataset']['train']['dataset']['args']['data_path'][0]

img_mode = config['dataset']['train']['dataset']['args']['img_mode']

pre_processes = config['dataset']['train']['dataset']['args']['pre_processes']

transform = config['dataset']['train']['dataset']['args']['transforms']

batch_size = config['dataset']['train']['loader']['batch_size']

shuffle = config['dataset']['train']['loader']['shuffle']

pin_memory = config['dataset']['train']['loader']['pin_memory']

num_workers = config['dataset']['train']['loader']['num_workers']

train_set = Dataset(data_path=data_path, img_mode=img_mode, pre_processes=pre_processes,

filter_keys=['img_path', 'img_name', 'text_polys', 'texts', 'ignore_tags'],

ignore_tags=['*', '###'], mode='train', transform=get_transforms(transform))

train_sampler = None

if config['distributed']:

from torch.utils.data.distributed import DistributedSampler

train_sampler = DistributedSampler(train_set)

shuffle = False

pin_memory = True

train_loader = DataLoader(dataset=train_set, sampler=train_sampler, pin_memory=pin_memory,

batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

# print("======config['dataset']['validate']", config['dataset']['validate'])

data_path = config['dataset']['validate']['dataset']['args']['data_path'][0]

img_mode = config['dataset']['validate']['dataset']['args']['img_mode']

min_scale = config['dataset']['validate']['dataset']['args']['pre_processes']['args']['min_scale']

max_scale = config['dataset']['validate']['dataset']['args']['pre_processes']['args']['max_scale']

transform = config['dataset']['validate']['dataset']['args']['transforms']

batch_size = config['dataset']['validate']['loader']['batch_size']

shuffle = config['dataset']['validate']['loader']['shuffle']

pin_memory = config['dataset']['validate']['loader']['pin_memory']

num_workers = config['dataset']['validate']['loader']['num_workers']

eval_set = Dataset(data_path=data_path, img_mode=img_mode, pre_processes=None,

filter_keys=['img_path', 'img_name', 'text_polys', 'texts', 'ignore_tags'],

ignore_tags=['*', '###'], mode='eval',

min_scale=min_scale, max_scale=max_scale, transform=get_transforms(transform))

eval_sampler = None

if config['distributed']:

from torch.utils.data.distributed import DistributedSampler

eval_sampler = DistributedSampler(eval_set)

shuffle = False

pin_memory = True

eval_loader = DataLoader(dataset=eval_set, sampler=eval_sampler, pin_memory=pin_memory,

batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

from collections import OrderedDict

modelWeights = OrderedDict()

for k, v in weights.items():

name = 'module.' + k # add `module.`

modelWeights[name] = v

"""

weights = torch.load(checkpoint_path)

if distributed:

module_weights = weights_add_module(weights['state_dict'])

model.load_state_dict(module_weights)

logging.info('====载入预训练模型成功====')

else:

model.load_state_dict(weights['state_dict'])

logging.info('====载入预训练模型成功====')

if optimizer is not None:

state_dict = model.module.state_dict() if distributed else model.state_dict()

# 生成后面要继续训练的模型

state = {

'state_dict': state_dict,

'optimizer': optimizer.state_dict()

}

torch.save(state, model_path)

os.environ["CUDA_VISIBLE_DEVICES"] = '4'

args = init_args()

config = anyconfig.load(open(args.config_file, 'rb'))

# print('===config:', config)

if 'base' in config:

config = parse_config(config)

print('===config:', config)

if torch.cuda.device_count() > 1:

torch.cuda.set_device(args.local_rank)

torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=torch.cuda.device_count(),

rank=args.local_rank)

config['distributed'] = True

else:

config['distributed'] = False

config['local_rank'] = args.local_rank

logging.info(config['dataset']['train'])

student_model = build_model(config['arch']['type'], **config['arch'])

# print('==student_model:', student_model)

teacher_model = build_model(config['kd']['type'], **config['kd'])

print('==teacher_model:', teacher_model)

#

student_criterion = build_loss(config['loss'].pop('type'), **config['loss']).cuda()

teacher_criterion = build_kd_loss(config['kd']['loss'].pop('type'), **config['kd']['loss']).cuda()

post_process = get_post_processing(config['post_processing'])

train_loader = get_trainloader(dataset.ICDAR2015Dataset, config)

eval_loader = get_evalloader(dataset.ICDAR2015Dataset, config)

student_model = student_model.cuda()

teacher_model = teacher_model.cuda()

if config['distributed']:

student_model = nn.parallel.DistributedDataParallel(student_model, device_ids=[args.local_rank],

output_device=args.local_rank, broadcast_buffers=False,

find_unused_parameters=True)

teacher_model = nn.parallel.DistributedDataParallel(teacher_model, device_ids=[args.local_rank],

output_device=args.local_rank, broadcast_buffers=False,

find_unused_parameters=True)

student_checkpoint_path = config['train']['resume_checkpoint']

teacher_checkpoint_path = config['kd']['resume_checkpoint']

output_path = config['train']['output_path']

optimizer = optim.Adam(student_model.parameters(), lr=0.001, weight_decay=0.00005)

# load_weights(model, optimizer, config['distributed'], checkpoint_path='/red_detection/DBNet/code_pretrain_model/model_latest_express_code_7_13.pth')

load_weights(student_model, optimizer, config['distributed'], checkpoint_path=student_checkpoint_path)

load_weights(teacher_model, None, config['distributed'], checkpoint_path=teacher_checkpoint_path)

epochs = config['train']['epochs']

warmup_iters = config['lr_scheduler']['args']['warmup_epoch'] * len(train_loader)

# scheduler = WarmupPolyLR(optimizer, max_iters=epochs * len(train_loader),

# warmup_iters=warmup_iters, **config['lr_scheduler']['args'])

train(student_model, teacher_model, optimizer, epochs, student_criterion, teacher_criterion, train_loader, config,

post_process, eval_loader, output_path)

from matplotlib import pyplot as plt

plt.plot(lr_list)

import sys

os.environ["CUDA_VISIBLE_DEVICES"] = '3'

args = init_args()

assert os.path.exists(args.config_file)

config = anyconfig.load(open(args.config_file, 'rb'))

# print('===config:', config)

if 'base' in config:

config = parse_config(config)

print('===config:', config)

print('==torch.cuda.device_count():', torch.cuda.device_count())

if torch.cuda.device_count() > 1:

torch.cuda.set_device(args.local_rank)

torch.distributed.init_process_group(backend="nccl", init_method="env://", world_size=torch.cuda.device_count(),

rank=args.local_rank)

config['distributed'] = True

else:

config['distributed'] = False

config['local_rank'] = args.local_rank

train_loader = get_trainloader(dataset.ICDAR2015Dataset, config)

# eval_loader = get_evalloader(dataset.ICDAR2015Dataset, config)

output_path = './查看图片_dataloader'

if not os.path.exists(output_path):

os.mkdir(output_path)

epochs = 1

for epoch in range(epochs):

for i, data_info in enumerate(tqdm(train_loader)):

# if i < 1:

print('===data_info:', data_info.keys())

batch_img = data_info['img']

shrink_map = data_info['shrink_map']

# threshold_label = data_info['threshold_map']

batch_gt = data_info['gt']

print('== batch_img.shape', batch_img.shape)

print('===shrink_map.shape', shrink_map.shape)

# print(batch_img.shape, threshold_label.shape, threshold_label.shape, batch_gt.shape, data_shape)

for j in range(batch_img.shape[0]):

img = batch_img[j].numpy().transpose(1, 2, 0)

gt = batch_gt[j].numpy() * 255.

# print('===img.shape:', img.shape)

# shrink_label = shrink_map[j].numpy()*255.

gt = np.expand_dims(gt, axis=-1)

img = (img * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])) * 255.

img = np.clip(gt + img, 0, 255)

cv2.imwrite(os.path.join(output_path, str(i) + '_' + str(j) + '.jpg'), img[..., ::-1])

x = torch.rand((8, 3, 640, 640)).cuda()

model_config = {

'backbone': {'type': 'resnet18', 'pretrained': True, "in_channels": 3},

'neck': {'type': 'FPN', 'inner_channels': 256}, # 分割头，FPN or FPEM_FFM

'head': {'type': 'DBHead', 'out_channels': 2, 'k': 50},

}

model = build_model('Model', **model_config).cuda()

print(model)

y = model(x)

print('y.shape:', y.shape)