def main(model_path, backbone, scale, path, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
save_img_folder = os.path.join(save_path, 'img')
if not os.path.exists(save_img_folder):
os.makedirs(save_img_folder)
save_txt_folder = os.path.join(save_path, 'result')
if not os.path.exists(save_txt_folder):
os.makedirs(save_txt_folder)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone=backbone, pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=scale, gpu_id=gpu_id)
total_frame = 0.0
total_time = 0.0
for img_path in tqdm(img_paths):
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_txt_folder, 'res_' + img_name + '.txt')
_, boxes_list, t = model.predict(img_path)
total_frame += 1
total_time += t
# img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
# cv2.imwrite(os.path.join(save_img_folder, '{}.jpg'.format(img_name)), img)
np.savetxt(save_name,
boxes_list.reshape(-1, 8),
delimiter=',',
fmt='%d')
print('fps:{}'.format(total_frame / total_time))
return save_txt_folder
def main(model_path, path, save_path, gpu_id):
if os.path.exists(save_path):
shutil.rmtree(save_path, ignore_errors=True)
if not os.path.exists(save_path):
os.makedirs(save_path)
img_paths = [os.path.join(path, x) for x in os.listdir(path)]
net = PSENet(backbone='resnet152', pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=1, gpu_id=gpu_id)
pbar = tqdm(total=len(img_paths))
for img_path in img_paths:
img_name = os.path.basename(img_path).split('.')[0]
save_name = os.path.join(save_path, 'res_' + img_name + '.txt')
_, preds = model.predict(img_path)
np.savetxt(save_name, preds.reshape(-1, 8), delimiter=',', fmt='%d')
pbar.update(1)
pbar.close()
def main():
if config.output_dir is None:
config.output_dir = 'output'
if config.restart_training:
shutil.rmtree(config.output_dir, ignore_errors=True)
if not os.path.exists(config.output_dir):
os.makedirs(config.output_dir)
logger = setup_logger(os.path.join(config.output_dir, 'train_log'))
logger.info(config.print())
torch.manual_seed(config.seed) # 为CPU设置随机种子
if config.gpu_id is not None and torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
logger.info('train with gpu {} and pytorch {}'.format(
config.gpu_id, torch.__version__))
device = torch.device("cuda:0")
torch.cuda.manual_seed(config.seed) # 为当前GPU设置随机种子
torch.cuda.manual_seed_all(config.seed) # 为所有GPU设置随机种子
else:
logger.info('train with cpu and pytorch {}'.format(torch.__version__))
device = torch.device("cpu")
train_data = MyDataset(config.trainroot,
data_shape=config.data_shape,
n=config.n,
m=config.m,
transform=transforms.ToTensor())
train_loader = Data.DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=int(config.workers))
writer = SummaryWriter(config.output_dir)
model = PSENet(backbone=config.backbone,
pretrained=config.pretrained,
result_num=config.n)
if not config.pretrained and not config.restart_training:
model.apply(weights_init)
num_gpus = torch.cuda.device_count()
if num_gpus > 1:
model = nn.DataParallel(model)
model = model.to(device)
# dummy_input = torch.autograd.Variable(torch.Tensor(1, 3, 600, 800).to(device))
# writer.add_graph(model=model, input_to_model=dummy_input)
criterion = PSELoss(Lambda=config.Lambda,
ratio=config.OHEM_ratio,
reduction='mean')
# optimizer = torch.optim.SGD(model.parameters(), lr=config.lr, momentum=0.99)
optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
if config.checkpoint != '' and not config.restart_training:
start_epoch = load_checkpoint(config.checkpoint, model, logger, device)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
config.lr_decay_step,
gamma=config.lr_gamma,
last_epoch=start_epoch)
else:
start_epoch = config.start_epoch
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
config.lr_decay_step,
gamma=config.lr_gamma)
all_step = len(train_loader)
logger.info('train dataset has {} samples,{} in dataloader'.format(
train_data.__len__(), all_step))
epoch = 0
best_model = {'recall': 0, 'precision': 0, 'f1': 0, 'model': ''}
try:
for epoch in range(start_epoch + 1, config.epochs):
start = time.time()
train_loss, lr = train_epoch(model, optimizer, scheduler,
train_loader, device, criterion,
epoch, all_step, writer, logger)
logger.info(
'[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}'.format(
epoch, config.epochs, train_loss,
time.time() - start, lr))
# net_save_path = '{}/PSENet_{}_loss{:.6f}.pth'.format(config.output_dir, epoch,
# train_loss)
# save_checkpoint(net_save_path, model, optimizer, epoch, logger)
if (0.3 < train_loss < 0.4 and epoch % 4 == 0) or train_loss < 0.3:
recall, precision, f1 = eval(
model, os.path.join(config.output_dir, 'output'),
config.testroot, device)
logger.info(
'test: recall: {:.6f}, precision: {:.6f}, f1: {:.6f}'.
format(recall, precision, f1))
net_save_path = '{}/PSENet_{}_loss{:.6f}_r{:.6f}_p{:.6f}_f1{:.6f}.pth'.format(
config.output_dir, epoch, train_loss, recall, precision,
f1)
save_checkpoint(net_save_path, model, optimizer, epoch, logger)
if f1 > best_model['f1']:
best_model['recall'] = recall
best_model['precision'] = precision
best_model['f1'] = f1
best_model['model'] = net_save_path
writer.add_scalar(tag='Test/recall',
scalar_value=recall,
global_step=epoch)
writer.add_scalar(tag='Test/precision',
scalar_value=precision,
global_step=epoch)
writer.add_scalar(tag='Test/f1',
scalar_value=f1,
global_step=epoch)
writer.close()
except KeyboardInterrupt:
save_checkpoint('{}/final.pth'.format(config.output_dir), model,
optimizer, epoch, logger)
finally:
if best_model['model']:
shutil.copy(
best_model['model'],
'{}/best_r{:.6f}_p{:.6f}_f1{:.6f}.pth'.format(
config.output_dir, best_model['recall'],
best_model['precision'], best_model['f1']))
logger.info(best_model)
import config
from model import PSENet
import matplotlib.pyplot as plt
from utils.utils import show_img, draw_bbox
os.environ['CUDA_VISIBLE_DEVICES'] = str('2')
model_path = 'output/psenet_icd2015_resnet152_my_loss_0.0001_author_crop_adam_newcrop_authorloss_fxw/PSENet_599_loss0.110923.pth'
# model_path = 'output/psenet_icd2015_new_loss/final.pth'
img_id = 1
img_path = '/data2/dataset/ICD15/test/img/img_{}.jpg'.format(img_id)
# img_path = '0.jpg'
label_path = '/data2/dataset/ICD15/test/gt/gt_img_{}.txt'.format(img_id)
label = _get_annotation(label_path)
img_path = '/data1/gcz/拍照清单数据集_备份/87436979.jpg'
# 初始化网络
net = PSENet(backbone='resnet152', pretrained=False, result_num=config.n)
model = Pytorch_model(model_path, net=net, scale=1, gpu_id=0)
# for i in range(100):
# model.predict(img_path)
preds, boxes_list = model.predict(img_path)
show_img(preds)
img = draw_bbox(img_path, boxes_list, color=(0, 0, 255))
cv2.imwrite('result.jpg', img)
# img = draw_bbox(img, label,color=(0,0,255))
show_img(img, color=True)
plt.show()
from data_generator import pasi_data
from loss_metric import score_loss, siam_loss, locate_loss, score_metric, locate_metric
from self_callbacks import MyEarlyStop
from acc_opt import SGDAccumulate
global patient_dict
os.environ["CUDA_VISIBLE_DEVICES"] = myModelConfig.availiable_gpus
steps_per_epoch_train = int(myModelConfig.num_train_examples_per_epoch // myModelConfig.batch_size)
steps_per_epoch_val = int(myModelConfig.num_val_examples_per_epoch // myModelConfig.batch_size)
data_loader = pasi_data()
train_gen = data_loader.train_generator(myModelConfig.batch_size)
valid_gen = data_loader.valid_generator(myModelConfig.batch_size)
siamese_model = PSENet(myModelConfig)
print(siamese_model.input)
print(siamese_model.output)
siamese_model.summary()
# parallel_model = siamese_model
parallel_model = multi_gpu_model(siamese_model, gpus=myModelConfig.num_gpus)
sgd_accu = SGDAccumulate(lr=myModelConfig.learning_rate, momentum=myModelConfig.momentum, nesterov=True,
accum_iters=myModelConfig.accu_num)
# sgd = optimizers.SGD(lr=learning_rate, momentum=momentum, nesterov=True)
reduce_lr = callbacks.ReduceLROnPlateau(monitor='val_loss', factor=myModelConfig.learning_rate_decay_factor, patience=5,
verbose=1, mode='min', cooldown=10, min_lr=0.00001)
my_early_stop = MyEarlyStop(siamese_model, myModelConfig.checkpoint_dir)