def test():
# get device
device = get_device(0)
# load net
num_classes = len(VOC_CLASSES)
testset = VOCDetection(args.voc_root, [('2007', 'test')], None, VOCAnnotationTransform())
mean = config.MEANS
cfg = config.voc_ab
if args.version == 'yolo_v2':
net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE)
print('Let us test yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test_net(net, device, testset,
BaseTransform(net.input_size, mean),
thresh=args.visual_threshold)
def test():
# get device
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# load net
num_classes = len(VOC_CLASSES)
testset = VOCDetection(args.voc_root, [('2007', 'test')], None,
VOCAnnotationTransform())
cfg = config.voc_ab
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
net = myYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.ANCHOR_SIZE)
print('Let us test yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.MULTI_ANCHOR_SIZE)
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
net = SlimYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.ANCHOR_SIZE)
print('Let us test slim-yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
net = YOLOv3tiny(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.TINY_MULTI_ANCHOR_SIZE)
print('Let us test tiny-yolo-v3 on the VOC0712 dataset ......')
net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test_net(net,
device,
testset,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
thresh=args.visual_threshold)
def test():
# get device
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# load net
num_classes = 80
if args.dataset == 'COCO':
cfg = config.coco_ab
testset = COCODataset(
data_dir=args.dataset_root,
json_file='instances_val2017.json',
name='val2017',
img_size=cfg['min_dim'][0],
debug=args.debug)
elif args.dataset == 'VOC':
cfg = config.voc_ab
testset = VOCDetection(VOC_ROOT, [('2007', 'test')], None, VOCAnnotationTransform())
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, anchor_size=config.ANCHOR_SIZE_COCO)
print('Let us test yolo-v2 on the MSCOCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, anchor_size=config.MULTI_ANCHOR_SIZE_COCO)
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
net = SlimYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, anchor_size=config.ANCHOR_SIZE_COCO)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
net = YOLOv3tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, anchor_size=config.TINY_MULTI_ANCHOR_SIZE_COCO)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test_net(net, device, testset,
BaseTransform(net.input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)),
thresh=args.visual_threshold)
def run():
args = parse_args()
# use cuda
if args.cuda:
device = torch.device("cuda")
else:
device = torch.device("cpu")
input_size = [args.input_size, args.input_size]
# load net
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
anchor_size = config.ANCHOR_SIZE_COCO
net = myYOLOv2(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size)
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
anchor_size = config.MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size)
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
anchor_size = config.ANCHOR_SIZE_COCO
net = SlimYOLOv2(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
anchor_size = config.TINY_MULTI_ANCHOR_SIZE_COCO
net = YOLOv3tiny(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# run
if args.mode == 'camera':
detect(net, device, BaseTransform(net.input_size),
thresh=args.visual_threshold, mode=args.mode)
elif args.mode == 'image':
detect(net, device, BaseTransform(net.input_size),
thresh=args.visual_threshold, mode=args.mode, path_to_img=args.path_to_img)
elif args.mode == 'video':
detect(net, device, BaseTransform(net.input_size),
thresh=args.visual_threshold, mode=args.mode, path_to_vid=args.path_to_vid, path_to_save=args.path_to_saveVid)
def test():
# get device
device = get_device(0)
# load net
num_classes = 80
anchor_size = config.ANCHOR_SIZE_COCO
if args.dataset == 'COCO':
cfg = config.coco_ab
testset = COCODataset(
data_dir=args.dataset_root,
json_file='instances_val2017.json',
name='val2017',
img_size=cfg['min_dim'][0],
debug=args.debug)
mean = config.MEANS
elif args.dataset == 'VOC':
cfg = config.voc_ab
testset = VOCDetection(VOC_ROOT, [('2007', 'test')], None, VOCAnnotationTransform())
mean = config.MEANS
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size)
print('Let us test yolo-v2 on the MSCOCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
net = YOLOv2tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
net = YOLOv3tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test_net(net, device, testset,
BaseTransform(net.input_size, mean),
thresh=args.visual_threshold)
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
num_classes = len(labelmap)
device = get_device(args.gpu_ind)
cfg = config.voc_ab
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE)
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE)
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
net = YOLOv2tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE)
print('Let us eval tiny-yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
net = YOLOv3tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE)
print('Let us eval tiny-yolo-v3 on the VOC0712 dataset ......')
# load net
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.eval()
print('Finished loading model!')
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
model = myYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
anchor_size=ANCHOR_SIZE_COCO)
print('Let us test yolo-v2 on the MSCOCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
model = myYOLOv3(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
anchor_size=MULTI_ANCHOR_SIZE_COCO)
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
model = YOLOv2tiny(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
anchor_size=ANCHOR_SIZE_COCO)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
model = YOLOv3tiny(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
if args.high_resolution == 1:
hr = True
cfg = coco_ab
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
total_anchor_size = tools.get_total_anchor_size(name='COCO')
yolo_net = myYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v2 on the MSCOCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
total_anchor_size = tools.get_total_anchor_size(multi_scale=True,
name='COCO')
yolo_net = myYOLOv3(device,
input_size=cfg['min_dim'],
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v3 on the MSCOCO dataset ......')
train(yolo_net, device)
def train():
args = parse_args().parse_args(args=[])
path_to_save = os.path.join(args.save_folder, args.dataset, args.version)
os.makedirs(path_to_save, exist_ok=True)
# use hi-res backbone
if args.high_resolution:
print('use hi-res backbone')
hr = True
else:
hr = False
# cuda
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda", 0)
else:
device = torch.device("cpu")
# multi-scale
if args.multi_scale:
print('use the multi-scale trick ...')
train_size = config_0['NIH_pancreas_data_aimshape']
val_size = config_0['NIH_pancreas_data_aimshape']
else:
train_size = config_0['NIH_pancreas_data_aimshape']
val_size = config_0['NIH_pancreas_data_aimshape']
cfg = train_cfg
# dataset and evaluator
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the dataset...')
if args.dataset == 'pnens':
num_classes = 1
# dataset_pkl = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl')
dataset_pkl = load_from_pkl(
r'E:\ly\pnens_data\nas_data\v1_data\NIH\pre_order0_128_128_64_new.pkl'
)
dataset = datanih(dataset_pkl)
evaluator = myEvaluator(dataset=dataset,
data_root="/data/data4T/ly/data/pnens_3D",
img_size=val_size,
device=device,
transform=BaseTransform(val_size),
labelmap=('pnens'))
else:
print('unknow dataset !! Only support voc and coco !!')
exit(0)
print('Training model on: yolov3_3D')
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# dataloader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True)
if args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
# anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
anchor_size = anchor_size_3D_try
yolo_net = myYOLOv3(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset))
else:
print('Unknown version !!!')
exit()
model = yolo_net
model.to(device).train()
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
# keep training
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(args.start_epoch, max_epoch):
# if epoch == 1: break
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * (
1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
"""
images [1, 1, 128, 128, 64]
targets [y1, y2, x1, x2, z1, z2, 0] 相对
"""
# targets [x1, x2, y1, y2, z1, z2, 有无物体(0 or 1)] 相对坐标[0~1]
# if iter_i == 0: break
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. /
(args.wp_epoch * epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
# to device
images = images.to(device)
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
size = random.randint(10, 19) * 32
train_size = [size, size]
model.set_grid(train_size)
if args.multi_scale:
# interpolate
# 上采样
images = torch.nn.functional.interpolate(images,
size=train_size,
mode='bilinear',
align_corners=False)
# make labels
# print(targets)
targets = [label.tolist() for label in targets]
targets = tools.multi_gt_creator3D(input_size=train_size,
strides=model.stride,
label_lists=targets,
anchor_size=anchor_size)
# [batch_index, featuremap_index， grid_x* grid_y * grid_z * ab_ind, 参数（15）]
# (15): [obj, class, tx, ty, tz, th, tw, td, weight, (xmin, ymin, zmin, ymax, xmax, zmax)(绝对)]
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss, dice_loss = model(
images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
# display
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(),
iter_i + epoch * epoch_size)
t1 = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f ||dice_loss %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), dice_loss.item(), cls_loss.item(),
txtytwth_loss.item(), total_loss.item(), train_size[0],
t1 - t0),
flush=True)
t0 = time.time()
# evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(val_size)
model.eval()
# evaluate
# evaluator.evaluate(model)
# convert to training mode.
model.trainable = True
model.set_grid(train_size)
model.train()
# save model
if (epoch + 1) % 5 == 0:
print('Saving state, saving in {} epoch{}:'.format(
os.path.join(path_to_save, args.version),
epoch + 1,
))
torch.save(
model.state_dict(),
os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth'))
else:
device = torch.device("cpu")
# input size
input_size = [args.input_size, args.input_size]
# load net
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
net = myYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'yolo_v3_spp':
from models.yolo_v3_spp import myYOLOv3Spp
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3Spp(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
net = SlimYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
anchor_size = TINY_MULTI_ANCHOR_SIZE if args.dataset == 'voc' else TINY_MULTI_ANCHOR_SIZE_COCO
net = YOLOv3tiny(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
def train():
args = parse_args()
path_to_save = os.path.join(args.save_folder, args.version)
os.makedirs(path_to_save, exist_ok=True)
hr = False
if args.high_resolution:
print('use hi-res backbone')
hr = True
cfg = voc_ab
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# use multi-scale trick
if args.multi_scale:
print('use multi-scale trick.')
input_size = [608, 608]
dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation([608, 608], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)))
else:
input_size = cfg['min_dim']
dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)))
# build model
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
total_anchor_size = tools.get_total_anchor_size()
yolo_net = myYOLOv2(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr)
print('Let us train yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
total_anchor_size = tools.get_total_anchor_size(multi_level=True, version='yolo_v3')
yolo_net = myYOLOv3(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr)
print('Let us train yolo-v3 on the VOC0712 dataset ......')
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
total_anchor_size = tools.get_total_anchor_size()
yolo_net = SlimYOLOv2(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr)
print('Let us train slim-yolo-v2 on the VOC0712 dataset ......')
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
total_anchor_size = tools.get_total_anchor_size(multi_level=True, version='tiny_yolo_v3')
yolo_net = YOLOv3tiny(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr)
print('Let us train tiny-yolo-v3 on the VOC0712 dataset ......')
else:
print('Unknown version !!!')
exit()
# finetune the model trained on COCO
if args.resume is not None:
print('finetune COCO trained ')
yolo_net.load_state_dict(torch.load(args.resume, map_location=device), strict=False)
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = os.path.join('log/voc/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
print("----------------------------------------Object Detection--------------------------------------------")
model = yolo_net
model.to(device)
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
# loss counters
print("----------------------------------------------------------")
print("Let's train OD network !")
print('Training on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
epoch_size = len(dataset) // args.batch_size
max_epoch = cfg['max_epoch']
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
t0 = time.time()
# start training
for epoch in range(max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5*(base_lr-0.00001)*(1+math.cos(math.pi*(epoch-20)*1./ (max_epoch-20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(data_loader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
targets = [label.tolist() for label in targets]
# make train label
if args.version == 'yolo_v2' or args.version == 'slim_yolo_v2':
targets = tools.gt_creator(input_size, yolo_net.stride, targets, version=args.version)
elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3':
targets = tools.multi_gt_creator(input_size, yolo_net.stride, targets, version=args.version)
# to device
images = images.to(device)
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss = model(images, target=targets)
# backprop and update
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size)
t1 = time.time()
print('[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch+1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), input_size[0], t1-t0),
flush=True)
t0 = time.time()
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
size = random.randint(10, 19) * 32
input_size = [size, size]
model.set_grid(input_size)
# change input dim
# But this operation will report bugs when we use more workers in data loader, so I have to use 0 workers.
# I don't know how to make it suit more workers, and I'm trying to solve this question.
data_loader.dataset.reset_transform(SSDAugmentation(input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)))
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(model.state_dict(), os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth')
)
def run():
args = parse_args()
if args.cuda:
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.setup == 'VOC':
print('use VOC style')
cfg = config.voc_ab
num_classes = 20
elif args.setup == 'COCO':
print('use COCO style')
cfg = config.coco_ab
num_classes = 80
else:
print('Only support VOC and COCO !!!')
exit(0)
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
if args.setup == 'VOC':
anchor_size = config.ANCHOR_SIZE
else:
anchor_size = config.ANCHOR_SIZE_COCO
net = myYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
trainable=False,
anchor_size=anchor_size)
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
if args.setup == 'VOC':
anchor_size = config.MULTI_ANCHOR_SIZE
else:
anchor_size = config.MULTI_ANCHOR_SIZE_COCO
net = myYOLOv3(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
trainable=False,
anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
if args.setup == 'VOC':
anchor_size = config.ANCHOR_SIZE
else:
anchor_size = config.ANCHOR_SIZE_COCO
net = YOLOv2tiny(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
trainable=False,
anchor_size=anchor_size)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
if args.setup == 'VOC':
anchor_size = config.MULTI_ANCHOR_SIZE
else:
anchor_size = config.MULTI_ANCHOR_SIZE_COCO
net = YOLOv3tiny(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
trainable=False,
anchor_size=anchor_size)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# run
if args.mode == 'image':
detect(net,
device,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
mode=args.mode,
path_to_img=args.path_to_img,
setup=args.setup)
elif args.mode == 'video':
detect(net,
device,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
mode=args.mode,
path_to_vid=args.path_to_vid,
path_to_save=args.path_to_saveVid,
setup=args.setup)
def train():
args = parse_args()
path_to_save = os.path.join(args.save_folder, args.dataset, args.version)
os.makedirs(path_to_save, exist_ok=True)
# use hi-res backbone
if args.high_resolution:
print('use hi-res backbone')
hr = True
else:
hr = False
# cuda
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# multi-scale
if args.multi_scale:
print('use the multi-scale trick ...')
train_size = [640, 640]
val_size = [416, 416]
else:
train_size = [416, 416]
val_size = [416, 416]
cfg = train_cfg
# dataset and evaluator
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the dataset...')
if args.dataset == 'voc':
data_dir = VOC_ROOT
num_classes = 20
dataset = VOCDetection(root=data_dir,
transform=SSDAugmentation(train_size))
evaluator = VOCAPIEvaluator(data_root=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size),
labelmap=VOC_CLASSES)
elif args.dataset == 'coco':
data_dir = coco_root
num_classes = 80
dataset = COCODataset(data_dir=data_dir,
img_size=train_size[0],
transform=SSDAugmentation(train_size),
debug=args.debug)
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size))
else:
print('unknow dataset !! Only support voc and coco !!')
exit(0)
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# dataloader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True)
# build model
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
yolo_net = myYOLOv2(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train yolo_v2 on the %s dataset ......' % (args.dataset))
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
yolo_net = myYOLOv3(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset))
elif args.version == 'yolo_v3_spp':
from models.yolo_v3_spp import myYOLOv3Spp
anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
yolo_net = myYOLOv3Spp(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train yolo_v3_spp on the %s dataset ......' %
(args.dataset))
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO
yolo_net = SlimYOLOv2(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train slim_yolo_v2 on the %s dataset ......' %
(args.dataset))
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
anchor_size = TINY_MULTI_ANCHOR_SIZE if args.dataset == 'voc' else TINY_MULTI_ANCHOR_SIZE_COCO
yolo_net = YOLOv3tiny(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train tiny_yolo_v3 on the %s dataset ......' %
(args.dataset))
else:
print('Unknown version !!!')
exit()
model = yolo_net
model.to(device).train()
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
# keep training
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(args.start_epoch, max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * (
1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. /
(args.wp_epoch * epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
# to device
images = images.to(device)
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
size = random.randint(10, 19) * 32
train_size = [size, size]
model.set_grid(train_size)
if args.multi_scale:
# interpolate
images = torch.nn.functional.interpolate(images,
size=train_size,
mode='bilinear',
align_corners=False)
# make labels
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2' or args.version == 'slim_yolo_v2':
targets = tools.gt_creator(input_size=train_size,
stride=yolo_net.stride,
label_lists=targets,
anchor_size=anchor_size)
else:
targets = tools.multi_gt_creator(input_size=train_size,
strides=yolo_net.stride,
label_lists=targets,
anchor_size=anchor_size)
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss = model(
images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
# display
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(),
iter_i + epoch * epoch_size)
t1 = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(),
total_loss.item(), train_size[0], t1 - t0),
flush=True)
t0 = time.time()
# evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(val_size)
model.eval()
# evaluate
evaluator.evaluate(model)
# convert to training mode.
model.trainable = True
model.set_grid(train_size)
model.train()
# save model
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
model.state_dict(),
os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth'))
img_size=input_size[0])
elif args.dataset == 'pnens':
num_classes = 1
class_names = ('pnens')
# dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl')
dataset2 = load_from_pkl('C:\\Users\\hazy\\Desktop\\fsdownload\\\pre_order0_128_128_64_new.pkl')
# dataset = rechange(dataset2)
dataset = datanih(dataset2)
class_colors = [(np.random.randint(255), np.random.randint(255), np.random.randint(255)) for _ in
range(num_classes)]
# load net
from models.yolo_v3 import myYOLOv3
anchor_size = anchor_size_3D_try
net = myYOLOv3(device, input_size=input_size, num_classes=1, conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh, anchor_size=anchor_size)
# net.load_state_dict(torch.load(args.trained_model, map_location=device))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test(net=net,
device=device,
testset=dataset,
transform=BaseTransform(input_size),
thresh=args.visual_threshold,
class_colors=class_colors,
class_names=class_names,
class_indexs=[0],
dataset_data_kind=args.dataset
def see_bbox():
args = parse_args().parse_args(args=[])
device = torch.device("cpu")
input_size = [128, 128, 64]
# dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl')
dataset2 = load_from_pkl(
r'E:\ly\pnens_data\nas_data\v1_data\NIH\pre_order0_128_128_64_new.pkl')
dataset = datanih(dataset2)
# load net
from models.yolo_v3 import myYOLOv3
anchor_size = anchor_size_3D_try
net = myYOLOv3(device,
input_size=input_size,
num_classes=1,
conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh,
anchor_size=anchor_size)
net.load_state_dict(
torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'),
map_location=device))
net.to(device).eval()
num_images = len(dataset)
result = []
images = []
masks = []
bbox_gt = []
for index in range(num_images):
if index == 1: break
# target
im, gt = dataset.__getitem__(index)
targets = get_targets(gt)
# gt [[obj, class, tx, ty, tz, tw, th, td, weight, (xmin, ymin, zmin, xmax, ymax, zmax)(绝对)]]
gt = targets[np.where(targets == 1)[0:2]]
txtytztwthtd = gt[0, 2:8]
gt_obj = targets[:, :, 0]
map_obj = obj_map_get(gt_obj)
print('Testing image {:d}/{:d}....'.format(index + 1, num_images))
img, _, mask, height, width, depth = dataset.pull_item(index)
bb = dataset.get_bbox_juedui(index) # [y1, y2, x1, x2, z1, z2] 绝对
bbox_gt.append(bb)
scale = np.array([[width, height, depth, width, height, depth]])
images.append(img)
masks.append(mask)
# to tensor
x = torch.tensor(img)
x = x.unsqueeze(0).to(device)
x = x.unsqueeze(0).to(device)
t0 = time.time()
# forward
bboxes, scores, _ = net(x) # bboxes[x1, y1, z1, x2, y2, z2] 相对
bboxes = bboxes * scale # to绝对
if not scores.size:
print('{} has not bbox'.format(index))
continue
best_scores = np.argmax(scores, axis=0)
bboxes = bboxes[best_scores]
# print(scores[best_scores])
# print([bboxes, best_scores])
result.append([bboxes, best_scores])
print("detection {} time used ".format(index), time.time() - t0, "s")
if False:
pred_1 = torch.sigmoid(pred_1)
pred_1 = pred_1.data.cpu().numpy()[0, 0]
img3D = pred_1
vol = mlab.pipeline.volume(mlab.pipeline.scalar_field(img3D),
name='3-d ultrasound ')
mlab.colorbar(orientation='vertical')
mlab.show()
show3Dslice(img3D)
img_resize = resize3D(img3D, [128, 64, 64])
show3Dslice(img_resize)
img = img.squeeze(0).data.cpu().numpy()
show3Dslice(img)
test_index = 0
# img3D[y, x, z]
img3D = images[test_index]
mask3D = masks[test_index]
# show3D(img3D)
# show3Dslice(img3D)
# 画bbox -- [x1, y1, z1, x2, y2, z2] 绝对
line_thick = 1
bbox3D = result[test_index][0]
x1, y1, z1, x2, y2, z2 = bbox3D
bbox3D = [y1, x1, z1, y2, x2, z2]
bbox3D = np.floor(bbox3D)
bbox3D = np.array(bbox3D, dtype=int)
# show3Dbbox_img(img3D, bbox3D, 2)
bboxInimg = bbox_in_img_for_slice(img3D, bbox3D, line_thick, -2e3)
# show3Dslice(bboxInimg)
bb_gt = bbox_gt[test_index] # [y1, y2, x1, x2, z1, z2] 绝对
y1, y2, x1, x2, z1, z2 = bb_gt
bb_gt = [y1, x1, z1, y2, x2, z2]
bb_gt = np.floor(bb_gt)
bb_gt = np.array(bb_gt, dtype=int)
print(bbox3D)
print(bb_gt)
bboxInimg = bbox_in_img_for_slice(bboxInimg, bb_gt, line_thick, 1e3)
bboxInimg = bboxInimg + mask3D * 2e3
# bbox_in_img = np.hstack((bbox_in_img, img3D))
show3Dslice(bboxInimg)
# show3D(bboxInimg)
img_for_3D = np.zeros([128, 128, 64], dtype=np.int)
img_for_3D = bbox_in_img_for_3D(img_for_3D, bb_gt, line_thick, 100)
img_for_3D = bbox_in_img_for_3D(img_for_3D, bbox3D, line_thick, 200)
show3D(img_for_3D)
def obj_test():
args = parse_args().parse_args(args=[])
device = torch.device("cpu")
input_size = [128, 128, 64]
# dataset
class_names = ('pnens')
# dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl')
dataset2 = load_from_pkl(r'E:\data\pre_order0_128_128_64_new.pkl')
dataset = datanih(dataset2)
# load net
from models.yolo_v3 import myYOLOv3
anchor_size = anchor_size_3D_try
net = myYOLOv3(device,
input_size=input_size,
num_classes=1,
conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh,
anchor_size=anchor_size)
net.load_state_dict(
torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'),
map_location=device))
net.to(device).eval()
num_images = len(dataset)
result = []
images = []
bbox_gt = []
for index in range(num_images):
if index == 4: break
print('Testing image {:d}/{:d}....'.format(index + 1, num_images))
img, _, height, width, depth = dataset.pull_item(index)
bb = dataset.get_bbox_juedui(index)
bbox_gt.append(bb)
scale = np.array([[height, width, depth, height, width, depth]])
images.append(img)
# to tensor
x = img
x = x.unsqueeze(0).to(device)
t0 = time.time()
# forward
bboxes, scores, _ = net(x)
bboxes = bboxes * scale
if not scores.size:
print('{} has not bbox'.format(index))
continue
best_scores = np.argmax(scores, axis=0)
bboxes = bboxes[best_scores]
# print(scores[best_scores])
# print([bboxes, best_scores])
result.append([bboxes, best_scores])
print("detection {} time used ".format(index), time.time() - t0, "s")
if False:
pred_1 = torch.sigmoid(pred_1)
pred_1 = pred_1.data.cpu().numpy()[0, 0]
img3D = pred_1
vol = mlab.pipeline.volume(mlab.pipeline.scalar_field(img3D),
name='3-d ultrasound ')
mlab.colorbar(orientation='vertical')
mlab.show()
show3Dslice(img3D)
img_resize = resize3D(img3D, [128, 64, 64])
show3Dslice(img_resize)
img = img.squeeze(0).data.cpu().numpy()
show3Dslice(img)
test_index = 1
img3D = images[test_index]
img3D = img3D.squeeze(dim=0)
img3D = np.array(img3D)
# show3D(img3D)
# show3Dslice(img3D)
# 画bbox -- [x1, y1, z1, x2, y2, z2]
line_thick = 1
bbox3D = result[test_index][0]
# print(bbox3D)
# print(result[0][1])
bbox3D = np.floor(bbox3D)
bbox3D = np.array(bbox3D, dtype=int)
# show3Dbbox_img(img3D, bbox3D, 2)
bbox_in_img = get_bbox_in_img(img3D, bbox3D, line_thick)
bb_gt = bbox_gt[test_index]
x1, x2, y1, y2, z1, z2 = bb_gt
bb_gt = [x1, y1, z1, x2, y2, z2]
bb_gt = np.floor(bb_gt)
bb_gt = np.array(bb_gt, dtype=int)
print(bbox3D)
print(bb_gt)
bbox_in_img = get_bbox_in_img(bbox_in_img,
bb_gt,
line_thick,
line_value=2e3)
# bbox_in_img = np.hstack((bbox_in_img, img3D))
show3Dslice(bbox_in_img)
# dataloader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True)
if args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
# anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO
anchor_size = anchor_size_3D_try
yolo_net = myYOLOv3(device,
input_size=train_size,
num_classes=num_classes,
trainable=True,
anchor_size=anchor_size,
hr=hr)
print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset))
else:
print('Unknown version !!!')
exit()
model = yolo_net
model.to(device).train()
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
def train():
args = parse_args()
data_dir = args.dataset_root
path_to_save = os.path.join(args.save_folder, args.version)
os.makedirs(path_to_save, exist_ok=True)
hr = False
if args.high_resolution:
print('use hi-res backbone')
hr = True
cfg = coco_ab
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.multi_scale:
print('Let us use the multi-scale trick.')
input_size = [608, 608]
dataset = COCODataset(data_dir=data_dir,
img_size=608,
transform=SSDAugmentation([608, 608],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)),
debug=args.debug)
else:
input_size = cfg['min_dim']
dataset = COCODataset(data_dir=data_dir,
img_size=cfg['min_dim'][0],
transform=SSDAugmentation(cfg['min_dim'],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)),
debug=args.debug)
# build model
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
total_anchor_size = tools.get_total_anchor_size(name='COCO')
yolo_net = myYOLOv2(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v2 on the COCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
total_anchor_size = tools.get_total_anchor_size(multi_level=True,
name='COCO')
yolo_net = myYOLOv3(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v3 on the COCO dataset ......')
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
total_anchor_size = tools.get_total_anchor_size(name='COCO')
yolo_net = YOLOv2tiny(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train tiny-yolo-v2 on the COCO dataset ......')
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
total_anchor_size = tools.get_total_anchor_size(multi_level=True,
name='COCO')
yolo_net = YOLOv3tiny(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train tiny-yolo-v3 on the COCO dataset ......')
else:
print('Unknown version !!!')
exit()
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
print('Let us train yolo-v2 on the MSCOCO dataset ......')
model = yolo_net
model.to(device).train()
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers)
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim'],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * (
1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(cfg['min_dim'])
# evaluate
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
# convert to training mode.
model.trainable = True
model.set_grid(input_size)
model.train()
if args.tfboard:
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
# tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4)
tmp_lr = 1e-6 + (base_lr - 1e-6) * (
iter_i + epoch * epoch_size) / (epoch_size *
(args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2' or args.version == 'tiny_yolo_v2':
targets = tools.gt_creator(input_size,
yolo_net.stride,
targets,
name='COCO')
elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3':
targets = tools.multi_gt_creator(input_size,
yolo_net.stride,
targets,
name='COCO')
# to device
images = images.to(device)
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss = model(
images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(),
iter_i + epoch * epoch_size)
if iter_i % 10 == 0:
t1 = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(),
total_loss.item(), input_size[0], t1 - t0),
flush=True)
t0 = time.time()
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# ms_ind = random.sample(ms_inds, 1)[0]
# input_size = cfg['multi_scale'][int(ms_ind)]
size = random.randint(10, 19) * 32
input_size = [size, size]
model.set_grid(input_size)
# change input dim
# But this operation will report bugs when we use more workers in data loader, so I have to use 0 workers.
# I don't know how to make it suit more workers, and I'm trying to solve this question.
dataloader.dataset.reset_transform(
SSDAugmentation(input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)))
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
model.state_dict(),
os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth'))
def see_feature_segment():
args = parse_args().parse_args(args=[])
device = torch.device("cpu")
input_size = [128, 128, 64]
# dataset
num_classes = 1
class_names = ('pnens')
# dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl')
dataset2 = load_from_pkl(r'E:\data\pre_order0_128_128_64_new.pkl')
dataset = datanih(dataset2)
# load net
from models.yolo_v3 import myYOLOv3
anchor_size = anchor_size_3D_try
net = myYOLOv3(device,
input_size=input_size,
num_classes=1,
conf_thresh=args.conf_thresh,
nms_thresh=args.nms_thresh,
anchor_size=anchor_size)
net.load_state_dict(
torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'),
map_location=device))
net.to(device).eval()
num_images = len(dataset)
result = []
images = []
for index in range(num_images):
if index == 1: break
print('Testing image {:d}/{:d}....'.format(index + 1, num_images))
img, _, height, width, depth = dataset.pull_item(index)
scale = np.array([[height, width, depth, height, width, depth]])
images.append(img)
# to tensor
x = img
x = x.unsqueeze(0).to(device)
t0 = time.time()
# forward
bboxes, scores, _ = net(x)
pred_1, pred_2, pred_3 = net.get_feature_map(x)
bboxes = bboxes * scale
if not scores.size:
print('{} has not bbox'.format(index))
continue
best_scores = np.argmax(scores, axis=0)
bboxes = bboxes[best_scores]
result.append([bboxes, scores])
print("detection {} time used ".format(index), time.time() - t0, "s")
img3D = images[0]
img3D = img3D.squeeze(dim=0)
img3D = np.array(img3D)
pred_1 = torch.sigmoid(pred_1)
pred_1 = pred_1.data.cpu().numpy()[0, 0]
show3D(pred_1)
img_resize = resize3D(pred_1, [128, 128, 64])
# show3Dslice(img_resize)
img_and_fearturemap = np.hstack((img_resize, img3D))
show3Dslice(img_and_fearturemap)
