def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
self.logger = logger
# 根据YoloV2和YoloV3使用不同的配置文件
if opt.model == 'Yolo2':
cfgfile = 'configs/yolo2-voc.cfg'
elif opt.model == 'Yolo3':
cfgfile = 'configs/yolo3-coco.cfg'
# 初始化detector
self.detector = Darknet(cfgfile, device=opt.device).to(opt.device)
print_network(self.detector, logger=logger)
# 在--load之前加载weights文件(可选)
if opt.weights:
utils.color_print('Load Yolo weights from %s.' % opt.weights, 3)
self.detector.load_weights(opt.weights)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, opt, logger=None):
super(Model, self).__init__(config, kwargs)
self.opt = opt
# cfgfile = 'yolo-voc.cfg'
# self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
#
# # replace the pre-trained head with a new one
# self.detector.roi_heads.box_predictor = FastRCNNPredictor(in_features, config.DATA.NUM_CLASSESS + 1)
self.detector = yolov4(inference=True,
n_classes=config.DATA.NUM_CLASSESS)
# """
# 预训练模型
# """
# pretrained_dict = torch.load('pretrained/yolov4.pth')
# self.detector.load_state_dict(pretrained_dict)
self.yolov4loss = Yolo_loss(device=opt.device, batch=opt.batch_size)
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.classifier = Classifier() #.cuda(device=opt.device)
#####################
# Init weights
#####################
# self.classifier.apply(weights_init)
print_network(self.classifier)
self.optimizer = optim.Adam(self.classifier.parameters(),
lr=opt.lr,
betas=(0.95, 0.999))
# load networks
if opt.load:
pretrained_path = opt.load
self.load_network(self.classifier, 'G', opt.which_epoch,
pretrained_path)
# if self.training:
# self.load_network(self.discriminitor, 'D', opt.which_epoch, pretrained_path)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
self.logger = logger
kargs = {}
if opt.scale:
min_size = opt.scale
max_size = int(min_size / 3 * 4)
kargs = {
'min_size': min_size,
'max_size': max_size,
}
kargs.update({'box_nms_thresh': nms_thresh})
# 定义backbone和Faster RCNN模型
if opt.backbone is None or opt.backbone.lower() in [
'res50', 'resnet50'
]:
# 默认是带fpn的resnet50
self.detector = fasterrcnn_resnet50_fpn(pretrained=False, **kargs)
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
elif opt.backbone.lower() in ['vgg16', 'vgg']:
backbone = vgg16_backbone()
self.detector = FasterRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res101', 'resnet101']:
# 不带FPN的resnet101
backbone = res101_backbone()
self.detector = FasterRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res', 'resnet']:
raise RuntimeError(
f'backbone "{opt.backbone}" is ambiguous, please specify layers.'
)
else:
raise NotImplementedError(f'no such backbone: {opt.backbone}')
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.classifier = Classifier() #.cuda(device=opt.device)
#####################
# Init weights
#####################
# self.classifier.apply(weights_init)
print_network(self.classifier)
self.optimizer = optim.Adam(self.classifier.parameters(),
lr=opt.lr,
betas=(0.95, 0.999))
# load networks
if opt.load:
pretrained_path = opt.load
self.load_network(self.classifier, 'G', opt.which_epoch,
pretrained_path)
# if self.training:
# self.load_network(self.discriminitor, 'D', opt.which_epoch, pretrained_path)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, input, label):
predicted = self.classifier(input)
loss = criterionCE(predicted, label)
self.avg_meters.update({'Cross Entropy': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {'predicted': predicted}
def forward(self, x):
return self.classifier(x)
def save(self, which_epoch):
self.save_network(self.classifier, 'G', which_epoch)
# self.save_network(self.discriminitor, 'D', which_epoch)
def update_learning_rate(self):
lrd = self.opt.lr / self.opt.niter_decay
lr = self.old_lr - lrd
# for param_group in self.d_optimizer.param_groups:
# param_group['lr'] = lr
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
if self.opt.verbose:
print('update learning rate: %f -> %f' % (self.old_lr, lr))
self.old_lr = lr
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.detector = get_net().to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, opt, logger=None):
super(Model, self).__init__(config, kwargs)
self.opt = opt
self.detector = get_net().to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=False)
# self.detector = FasterRCNN_VGG()
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
# 根据YoloV2和YoloV3使用不同的配置文件
if config.MODEL.NAME == 'Yolo2':
cfgfile = 'configs/networks/yolo2-voc.cfg'
elif config.MODEL.NAME == 'Yolo3':
cfgfile = 'configs/networks/yolo3-coco.cfg'
# 初始化detector
self.detector = Darknet(cfgfile, device=opt.device).to(opt.device)
if opt.debug:
print_network(self.detector)
# 在--load之前加载weights文件(可选)
if opt.load and opt.load[-2:] != 'pt':
if is_first_gpu():
utils.color_print('Load Yolo weights from %s.' % opt.load, 3)
self.detector.load_weights(opt.load)
elif 'LOAD' in config.MODEL and config.MODEL.LOAD[-2:] != 'pt':
if is_first_gpu():
utils.color_print(
'Load Yolo weights from %s.' % config.MODEL.LOAD, 3)
self.detector.load_weights(config.MODEL.LOAD)
self.to(opt.device)
# 多GPU支持
if is_distributed():
self.detector = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
self.detector)
self.detector = torch.nn.parallel.DistributedDataParallel(
self.detector,
find_unused_parameters=False,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# self.detector = torch.nn.parallel.DistributedDataParallel(self.detector, device_ids=[opt.local_rank], output_device=opt.local_rank)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
cfgfile = 'configs/yolov5x.yaml'
self.detector = Yolo5(cfgfile)
self.detector.hyp = hyp
self.detector.gr = 1.0
self.detector.nc = opt.num_classes
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
self.it = 0
def __init__(self, opt, logger=None):
super(Model, self).__init__(config, kwargs)
self.opt = opt
cfgfile = 'configs/yolov5x.yaml'
self.detector = Yolo5(cfgfile)
self.detector.hyp = hyp
self.detector.gr = 1.0
self.detector.nc = config.DATA.NUM_CLASSESS
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
self.it = 0
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.classifier = Classifier()
#####################
# Init weights
#####################
# self.classifier.apply(weights_init)
print_network(self.classifier)
self.optimizer = optim.Adam(self.classifier.parameters(),
lr=opt.lr,
betas=(0.95, 0.999))
self.scheduler = get_scheduler(opt, self.optimizer)
# load networks
# if opt.load:
# pretrained_path = opt.load
# self.load_network(self.classifier, 'G', opt.which_epoch, pretrained_path)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
# cfgfile = 'yolo-voc.cfg'
# self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
#
# # replace the pre-trained head with a new one
# self.detector.roi_heads.box_predictor = FastRCNNPredictor(in_features, opt.num_classes + 1)
self.detector = Retina_50(opt.num_classes,pretrained=True)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
self.detector = SSDDetector(opt).to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
CENTER_VARIANCE = 0.1
SIZE_VARIANCE = 0.2
THRESHOLD = 0.5
self.target_transform = SSDTargetTransform(
PriorBox(opt)(), CENTER_VARIANCE, SIZE_VARIANCE, THRESHOLD)
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
self.detector = SSDDetector(config).to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
CENTER_VARIANCE = 0.1
SIZE_VARIANCE = 0.2
THRESHOLD = 0.5
self.target_transform = SSDTargetTransform(
PriorBox(config)(), CENTER_VARIANCE, SIZE_VARIANCE, THRESHOLD)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
# 根据YoloV2和YoloV3使用不同的配置文件
if opt.model == 'Yolo2':
cfgfile = 'configs/yolo2-voc.cfg'
elif opt.model == 'Yolo3':
cfgfile = 'configs/yolo3-coco.cfg'
# 初始化detector
self.detector = Darknet(cfgfile, device=opt.device).to(opt.device)
print_network(self.detector)
# 在--load之前加载weights文件(可选)
if opt.weights:
utils.color_print('Load Yolo weights from %s.' % opt.weights, 3)
self.detector.load_weights(opt.weights)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *arg):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
loss_layers = self.detector.loss_layers
org_loss = []
image = sample['image'].to(opt.device) # target domain
target = sample['yolo_boxes'].to(opt.device)
detection_output = self.detector(image) # 在src domain上训练检测
for i, l in enumerate(loss_layers):
ol = l(detection_output[i]['x'], target)
org_loss.append(ol)
loss = sum(org_loss)
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.detector.parameters(), 10000)
self.optimizer.step()
org_loss.clear()
gc.collect()
return {}
def forward(self, image):
"""
Args:
image: [b, 3, h, w] Tensor
"""
batch_bboxes = []
batch_labels = []
batch_scores = []
if self.detector.net_name() == 'region': # region_layer
shape = (0, 0)
else:
shape = (opt.width, opt.height)
num_classes = self.detector.num_classes
output = self.detector(image)
all_boxes = get_all_boxes(output,
shape,
conf_thresh,
num_classes,
device=opt.device,
only_objectness=0,
validation=True)
for b in range(len(all_boxes)):
boxes = all_boxes[b]
width = opt.width
height = opt.height
correct_yolo_boxes(boxes, width, height, width, height)
boxes = nms(boxes, nms_thresh)
img_boxes = []
img_labels = []
img_scores = []
b, c, h, w = image.shape
for box in boxes:
box[0] -= box[2] / 2
box[1] -= box[3] / 2
box[2] += box[0]
box[3] += box[1]
box[0] *= w
box[2] *= w
box[1] *= h
box[3] *= h
for i in range(5, len(box), 2):
img_boxes.append(box[:4])
img_labels.append(box[i + 1])
score = box[4] * box[i]
img_scores.append(score)
batch_bboxes.append(np.array(img_boxes))
batch_labels.append(np.array(img_labels).astype(np.int32))
batch_scores.append(np.array(img_scores))
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.detector = SSDDetector(opt).to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
CENTER_VARIANCE = 0.1
SIZE_VARIANCE = 0.2
THRESHOLD = 0.5
self.target_transform = SSDTargetTransform(
PriorBox(opt)(), CENTER_VARIANCE, SIZE_VARIANCE, THRESHOLD)
def update(self, sample, *arg):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
loss_dict = self.forward(sample)
loss = sum(loss for loss in loss_dict.values())
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward(self, sample):
if self.training:
return self.forward_train(sample)
else:
return self.forward_test(sample)
def forward_train(self, sample):
labels = sample['labels']
for label in labels:
label += 1. # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
for box in bboxes:
box[:, 0::2] /= 512 # width
box[:, 1::2] /= 512 # height
for b in range(len(labels)):
bboxes[b], labels[b] = self.target_transform(bboxes[b], labels[b])
image = image.to(opt.device)
bboxes = torch.stack(bboxes).to(opt.device)
labels = torch.stack(labels).long().to(opt.device)
targets = {'boxes': bboxes, 'labels': labels}
loss_dict = self.detector(image, targets=targets)
return loss_dict
def forward_test(self, image):
conf_thresh = 0.5
batch_bboxes = []
batch_labels = []
batch_scores = []
outputs = self.detector(image)
for b in range(len(outputs)): #
output = outputs[b]
boxes = output['boxes']
labels = output['labels']
scores = output['scores']
boxes = boxes[scores > conf_thresh]
labels = labels[scores > conf_thresh]
labels = labels - 1
scores = scores[scores > conf_thresh]
batch_bboxes.append(boxes.detach().cpu().numpy())
batch_labels.append(labels.detach().cpu().numpy())
batch_scores.append(scores.detach().cpu().numpy())
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
if opt.scale:
min_size = opt.scale
max_size = int(min_size / 3 * 5)
else:
min_size = 800
max_size = 1333
# anchor_sizes = ((16,), (32,), (64,), (128,), (512,)) # ,( 4,), (256,), (512,))
# aspect_ratios = ((0.2, 0.5, 1.0, 2.0, 5.0),) * len(anchor_sizes)
# rpn_anchor_generator = AnchorGenerator(
# anchor_sizes, aspect_ratios
# )
kargs = {
'min_size': min_size,
'max_size': max_size,
'cascade_iou_thr': [0.5, 0.6, 0.7],
}
# 定义backbone和Faster RCNN模型
if opt.backbone is None or opt.backbone.lower() in [
'res50', 'resnet50'
]:
# 默认是带fpn的resnet50
self.detector = cascadercnn_resnet50_fpn(pretrained=False, **kargs)
in_features = self.detector.roi_heads[
0].box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads[0].box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
self.detector.roi_heads[1].box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
self.detector.roi_heads[2].box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
elif opt.backbone.lower() in ['vgg16', 'vgg']:
backbone = vgg16_backbone()
self.detector = CascadeRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res101', 'resnet101']:
# 不带FPN的resnet101
backbone = res101_backbone()
self.detector = CascadeRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res', 'resnet']:
raise RuntimeError(
f'backbone "{opt.backbone}" is ambiguous, please specify layers.'
)
else:
raise NotImplementedError(f'no such backbone: {opt.backbone}')
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
class Model(BaseModel):
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
# cfgfile = 'yolo-voc.cfg'
# self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
#
# # replace the pre-trained head with a new one
# self.detector.roi_heads.box_predictor = FastRCNNPredictor(in_features, opt.num_classes + 1)
self.detector = Retina_50(opt.num_classes,pretrained=True)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *arg):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
images = sample['image'].to(opt.device)
annotations = [torch.cat([box, label.unsqueeze(1)], dim=1).to(opt.device) for box, label in zip(sample['bboxes'],sample['labels'])]
inputs = images, annotations
loss = sum(self.detector(inputs))
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward_test(self, image): # test
"""给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用"""
conf_thresh = 0. # 0.5 for vis result
batch_bboxes = []
batch_labels = []
batch_scores = []
for i in range(image.shape[0]):
single_image = image[i: i+1]
scores, labels, bboxes = self.detector(single_image) # RetinaNet只支持单张检测
conf = scores > conf_thresh
bboxes = bboxes[conf].detach().cpu().numpy()
labels = labels[conf].detach().cpu().numpy()
scores = scores[conf].detach().cpu().numpy()
batch_bboxes.append(bboxes)
batch_labels.append(labels)
batch_scores.append(scores)
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
# 根据YoloV2和YoloV3使用不同的配置文件
if config.MODEL.NAME == 'Yolo2':
cfgfile = 'configs/networks/yolo2-voc.cfg'
elif config.MODEL.NAME == 'Yolo3':
cfgfile = 'configs/networks/yolo3-coco.cfg'
# 初始化detector
self.detector = Darknet(cfgfile, device=opt.device).to(opt.device)
if opt.debug:
print_network(self.detector)
# 在--load之前加载weights文件(可选)
if opt.load and opt.load[-2:] != 'pt':
if is_first_gpu():
utils.color_print('Load Yolo weights from %s.' % opt.load, 3)
self.detector.load_weights(opt.load)
elif 'LOAD' in config.MODEL and config.MODEL.LOAD[-2:] != 'pt':
if is_first_gpu():
utils.color_print(
'Load Yolo weights from %s.' % config.MODEL.LOAD, 3)
self.detector.load_weights(config.MODEL.LOAD)
self.to(opt.device)
# 多GPU支持
if is_distributed():
self.detector = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
self.detector)
self.detector = torch.nn.parallel.DistributedDataParallel(
self.detector,
find_unused_parameters=False,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# self.detector = torch.nn.parallel.DistributedDataParallel(self.detector, device_ids=[opt.local_rank], output_device=opt.local_rank)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def update(self, sample, *args):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
#loss_layers = self.detector.loss_layers
#org_loss = []
image = torch.stack(sample['image']).to(opt.device) # target domain
target = sample['yolo_boxes'].to(opt.device)
#detection_output = self.detector(image) # 在src domain上训练检测
#for i, l in enumerate(loss_layers):
# ol=l(detection_output[i]['x'], target)
# org_loss.append(ol)
#loss = sum(org_loss)
loss = self.detector(image, target=target)
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.detector.parameters(), 10000)
self.optimizer.step()
return {}
def forward_test(self, image):
"""
给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用
Args:
image: [b, 3, h, w] list
"""
image = torch.stack(image).to(opt.device)
batch_bboxes = []
batch_labels = []
batch_scores = []
if self.config.MODEL.NAME == 'Yolo2': # region_layer
shape = (0, 0)
else:
shape = self.config.DATA.SCALE
assert len(shape) == 2
num_classes = self.config.DATA.NUM_CLASSESS
outputs = self.detector(image)
outputs = get_all_boxes(outputs,
shape,
conf_thresh,
num_classes,
device=opt.device,
only_objectness=False,
validation=True)
# ipdb.set_trace()
for b in range(image.shape[0]):
preds = outputs[b]
preds = preds[preds[:, 4] > conf_thresh]
boxes = preds[:, :4]
det_conf = preds[:, 4]
cls_conf = preds[:, 5:]
if len(cls_conf):
_, labels = torch.max(cls_conf, 1)
b, c, h, w = image.shape
boxes = xywh_to_xyxy(boxes, w, h) # yolo的xywh转成输出的xyxy
nms_indices = box_ops.batched_nms(boxes, det_conf, labels,
nms_thresh)
# nms_indices = nms(boxes, det_conf, nms_thresh)
else:
nms_indices = []
if len(nms_indices) == 0:
batch_bboxes.append(np.array([[]], np.float32))
batch_labels.append(np.array([], np.int32))
batch_scores.append(np.array([], np.float32))
continue
boxes, det_conf, cls_conf = keep(nms_indices,
[boxes, det_conf, cls_conf])
max_conf, labels = torch.max(cls_conf, 1)
scores = det_conf * max_conf
if not DO_FAST_EVAL: # 只保留概率最高的类，能够加快eval速度但会降低精度
repeat_boxes = boxes.repeat([num_classes, 1])
repeat_det_conf = det_conf.repeat(num_classes)
repeat_conf = cls_conf.transpose(1, 0).contiguous().view(-1)
repeat_labels = torch.linspace(
0, num_classes - 1,
num_classes).repeat(boxes.shape[0],
1).transpose(1,
0).contiguous().view(-1)
boxes, det_conf, cls_conf, labels = keep(
repeat_conf > cls_thresh, [
repeat_boxes, repeat_det_conf, repeat_conf,
repeat_labels
])
scores = det_conf * cls_conf
batch_bboxes.append(to_numpy(boxes))
batch_labels.append(to_numpy(labels, np.int32))
batch_scores.append(to_numpy(scores))
return batch_bboxes, batch_labels, batch_scores
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
cfgfile = 'configs/yolov5x.yaml'
self.detector = Yolo5(cfgfile)
self.detector.hyp = hyp
self.detector.gr = 1.0
self.detector.nc = opt.num_classes
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
self.it = 0
def update(self, sample, *arg):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
# self.it += 1
# ni = self.it + self.nb * self.epoch
#
# if ni <= self.n_burn:
# xi = [0, self.n_burn] # x interp
# # model.gr = np.interp(ni, xi, [0.0, 1.0]) # giou loss ratio (obj_loss = 1.0 or giou)
# accumulate = max(1, np.interp(ni, xi, [1, 64 / opt.batch_size]).round())
# for j, x in enumerate(self.optimizer.param_groups):
# # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
# x['lr'] = np.interp(ni, xi, [0.1 if j == 2 else 0.0, x['initial_lr'] * (((1 + math.cos(self.epoch * math.pi / 300)) / 2) ** 1.0) * 0.9 + 0.1])
# if 'momentum' in x:
# x['momentum'] = np.interp(ni, xi, [0.9, hyp['momentum']])
image = sample['image'].to(opt.device) # target domain
target = sample['yolo5_boxes'].to(opt.device)
pred = self.detector(image)
loss, loss_items = compute_loss(pred, target.to(opt.device), self.detector)
self.avg_meters.update({'loss': sum(loss_items).item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward(self, image): # test
batch_bboxes = []
batch_labels = []
batch_scores = []
inf_out, _ = self.detector(image)
bboxes = non_max_suppression(inf_out, conf_thres=0.001, iou_thres=0.65, merge=False)
b = len(bboxes)
for bi in range(b):
pred = bboxes[bi]
if pred is None:
batch_bboxes.append(np.array([[]]))
batch_labels.append(np.array([]))
batch_scores.append(np.array([]))
else:
boxes = pred[:,:4].cpu().detach().numpy()
scores = pred[:,4].cpu().detach().numpy()
labels = pred[:,5].cpu().detach().numpy().astype(np.int32)
batch_bboxes.append(boxes)
batch_labels.append(labels)
batch_scores.append(scores)
# ipdb.set_trace()
"""xywh转x1y1x2y2"""
# bboxes[:, :, 0] = bboxes[:, :, 0] - bboxes[:, :, 2] / 2
# bboxes[:, :, 1] = bboxes[:, :, 1] - bboxes[:, :, 3] / 2
# bboxes[:, :, 2] += bboxes[:, :, 0]
# bboxes[:, :, 3] += bboxes[:, :, 1]
# b = image.shape[0] # batch有几张图
# for bi in range(b):
# bbox = bboxes[bi]
# conf_bbox = bbox[bbox[:, 4] > opt.conf_thresh]
# xyxy_bbox = conf_bbox[:, :4] # x1y1x2y2坐标
# scores = conf_bbox[:, 4]
# nms_indices = nms(xyxy_bbox, scores, opt.nms_thresh)
# xyxy_bbox = xyxy_bbox[nms_indices]
# scores = scores[nms_indices] # 检测的置信度
# classification = conf_bbox[nms_indices, 5:]
# if len(classification) != 0:
# prob, class_id = torch.max(classification, 1)
# # scores = scores * prob # 乘以最高类别的置信度
# else:
# class_id = torch.Tensor([])
# if opt.box_fusion == 'wbf':
# pass
# # boxes, scores, labels = weighted_boxes_fusion([xyxy_bbox.detach().cpu().numpy()],
# # [scores.detach().cpu().numpy()],
# # [np.zeros([xyxy_bbox.shape[0]], dtype=np.int32)],
# # iou_thr=0.5)
# elif opt.box_fusion == 'nms':
# boxes = xyxy_bbox.detach().cpu().numpy()
# scores = scores.detach().cpu().numpy()
# labels = class_id.detach().cpu().numpy()
# batch_bboxes.append(boxes)
# batch_labels.append(labels)
# batch_scores.append(scores)
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
# state = torch.load(ckpt_path)
# utils.p(list(state['detector'].keys()))
# print('=========================================')
# utils.p(list(self.detector.state_dict().keys()))
# ipdb.set_trace()
# self.detector.load_state_dict(state)
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.classifier = Classifier() #.cuda(device=opt.device)
#####################
# Init weights
#####################
# self.classifier.apply(weights_init)
print_network(self.classifier)
self.optimizer = get_optimizer(opt, self.classifier)
self.scheduler = get_scheduler(opt, self.optimizer)
# load networks
# if opt.load:
# pretrained_path = opt.load
# self.load_network(self.classifier, 'G', opt.which_epoch, pretrained_path)
# if self.training:
# self.load_network(self.discriminitor, 'D', opt.which_epoch, pretrained_path)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, input, label):
predicted = self.classifier(input)
loss_ce = self.criterionCE(predicted, label)
loss = loss_ce
self.avg_meters.update({'Cross Entropy': loss_ce.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {'predicted': predicted}
def forward(self, x):
return self.classifier(x)
def load(self, ckpt_path):
load_dict = torch.load(ckpt_path, map_location=opt.device)
self.classifier.load_state_dict(load_dict['classifier'])
if opt.resume:
self.optimizer.load_state_dict(load_dict['optimizer'])
self.scheduler.load_state_dict(load_dict['scheduler'])
epoch = load_dict['epoch']
utils.color_print('Load checkpoint from %s, resume training.' % ckpt_path, 3)
else:
epoch = load_dict['epoch']
utils.color_print('Load checkpoint from %s.' % ckpt_path, 3)
return epoch
def save(self, which_epoch):
# self.save_network(self.classifier, 'G', which_epoch)
save_filename = f'{which_epoch}_{opt.model}.pt'
save_path = os.path.join(self.save_dir, save_filename)
save_dict = OrderedDict()
save_dict['classifier'] = self.classifier.state_dict()
# save_dict['discriminitor'] = self.discriminitor.state_dict()
save_dict['optimizer'] = self.optimizer.state_dict()
save_dict['scheduler'] = self.scheduler.state_dict()
save_dict['epoch'] = which_epoch
torch.save(save_dict, save_path)
utils.color_print(f'Save checkpoint "{save_path}".', 3)
class Model(BaseModel):
def __init__(self, opt, logger=None):
super(Model, self).__init__()
self.opt = opt
self.logger = logger
kargs = {}
if opt.scale:
min_size = opt.scale
max_size = int(min_size / 3 * 4)
kargs = {
'min_size': min_size,
'max_size': max_size,
}
kargs.update({'box_nms_thresh': nms_thresh})
# 定义backbone和Faster RCNN模型
if opt.backbone is None or opt.backbone.lower() in [
'res50', 'resnet50'
]:
# 默认是带fpn的resnet50
self.detector = fasterrcnn_resnet50_fpn(pretrained=False, **kargs)
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
elif opt.backbone.lower() in ['vgg16', 'vgg']:
backbone = vgg16_backbone()
self.detector = FasterRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res101', 'resnet101']:
# 不带FPN的resnet101
backbone = res101_backbone()
self.detector = FasterRCNN(backbone,
num_classes=opt.num_classes + 1,
**kargs)
elif opt.backbone.lower() in ['res', 'resnet']:
raise RuntimeError(
f'backbone "{opt.backbone}" is ambiguous, please specify layers.'
)
else:
raise NotImplementedError(f'no such backbone: {opt.backbone}')
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *arg):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
labels = sample['labels']
for label in labels:
label += 1. # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
for b in range(len(image)):
if len(bboxes[b]) == 0: # 没有bbox，不更新参数
return {}
image = image.to(opt.device)
bboxes = [bbox.to(opt.device).float() for bbox in bboxes]
labels = [label.to(opt.device).float() for label in labels]
image = list(im for im in image)
b = len(bboxes)
target = [{
'boxes': bboxes[i],
'labels': labels[i].long()
} for i in range(b)]
"""
target['boxes'] = boxes
target['labels'] = labels
# target['masks'] = None
target['image_id'] = torch.tensor([index])
target['area'] = area
target['iscrowd'] = iscrowd
"""
loss_dict = self.detector(image, target)
loss = sum(l for l in loss_dict.values())
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward_test(self, image): # test
"""给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用"""
image = list(im for im in image)
batch_bboxes = []
batch_labels = []
batch_scores = []
with torch.no_grad():
outputs = self.detector(image)
for b in range(len(outputs)): #
output = outputs[b]
boxes = output['boxes']
labels = output['labels']
scores = output['scores']
boxes = boxes[scores > conf_thresh]
labels = labels[scores > conf_thresh]
labels = labels.detach().cpu().numpy()
# for i in range(len(labels)):
# labels[i] = coco_90_to_80_classes(labels[i])
labels = labels - 1
scores = scores[scores > conf_thresh]
batch_bboxes.append(boxes.detach().cpu().numpy())
batch_labels.append(labels)
batch_scores.append(scores.detach().cpu().numpy())
return batch_bboxes, batch_labels, batch_scores
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
kargs = {}
if 'SCALE' in config.DATA:
scale = config.DATA.SCALE
if isinstance(scale, int):
min_size = scale
max_size = int(min_size / 3 * 5)
else:
min_size, max_size = config.DATA.SCALE
kargs = {
'min_size': min_size,
'max_size': max_size,
}
kargs.update({'box_nms_thresh': config.TEST.NMS_THRESH})
# 多卡使用 SyncBN
if is_distributed():
kargs.update({'norm_layer': torch.nn.SyncBatchNorm})
# 定义backbone和Faster RCNN模型
if config.MODEL.BACKBONE is None or config.MODEL.BACKBONE.lower() in [
'res50', 'resnet50'
]:
# 默认是带fpn的resnet50
self.detector = fasterrcnn_resnet50_fpn(pretrained=False, **kargs)
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, config.DATA.NUM_CLASSESS + 1)
elif config.MODEL.BACKBONE.lower() in ['vgg16', 'vgg']:
backbone = vgg16_backbone()
self.detector = FasterRCNN(backbone,
num_classes=config.DATA.NUM_CLASSESS +
1,
**kargs)
elif config.MODEL.BACKBONE.lower() in ['res101', 'resnet101']:
# 不带FPN的resnet101
backbone = res101_backbone()
self.detector = FasterRCNN(backbone,
num_classes=config.DATA.NUM_CLASSESS +
1,
**kargs)
elif config.MODEL.BACKBONE.lower() in ['res', 'resnet']:
raise RuntimeError(
f'backbone "{config.MODEL.BACKBONE}" is ambiguous, please specify layers.'
)
else:
raise NotImplementedError(
f'no such backbone: {config.MODEL.BACKBONE}')
if opt.debug and is_first_gpu():
print_network(self.detector)
self.to(opt.device)
# 多GPU支持
if is_distributed():
self.detector = torch.nn.parallel.DistributedDataParallel(
self.detector,
find_unused_parameters=False,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# self.detector = torch.nn.parallel.DistributedDataParallel(self.detector, device_ids=[opt.local_rank], output_device=opt.local_rank)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def update(self, sample, *arg):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
labels = sample['labels']
for label in labels:
label += 1. # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
for b in range(len(image)):
if len(bboxes[b]) == 0: # 没有bbox，不更新参数
return {}
#image = image.to(opt.device)
bboxes = [bbox.to(opt.device).float() for bbox in bboxes]
labels = [label.to(opt.device).float() for label in labels]
image = list(im.to(opt.device) for im in image)
b = len(bboxes)
target = [{
'boxes': bboxes[i],
'labels': labels[i].long()
} for i in range(b)]
"""
target['boxes'] = boxes
target['labels'] = labels
# target['masks'] = None
target['image_id'] = torch.tensor([index])
target['area'] = area
target['iscrowd'] = iscrowd
"""
loss_dict = self.detector(image, target)
loss = sum(l for l in loss_dict.values())
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward_test(self, image): # test
"""给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用"""
#image = list(im for im in image)
image = list(im.to(opt.device) for im in image)
batch_bboxes = []
batch_labels = []
batch_scores = []
with torch.no_grad():
outputs = self.detector(image)
conf_thresh = self.config.TEST.CONF_THRESH
for b in range(len(outputs)): #
output = outputs[b]
boxes = output['boxes']
labels = output['labels']
scores = output['scores']
boxes = boxes[scores > conf_thresh]
labels = labels[scores > conf_thresh]
labels = labels.detach().cpu().numpy()
# for i in range(len(labels)):
# labels[i] = coco_90_to_80_classes(labels[i])
labels = labels - 1
scores = scores[scores > conf_thresh]
batch_bboxes.append(boxes.detach().cpu().numpy())
batch_labels.append(labels)
batch_scores.append(scores.detach().cpu().numpy())
return batch_bboxes, batch_labels, batch_scores
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, opt, logger=None):
super(Model, self).__init__(config, kwargs)
self.opt = opt
# cfgfile = 'yolo-voc.cfg'
# self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
#
# # replace the pre-trained head with a new one
# self.detector.roi_heads.box_predictor = FastRCNNPredictor(in_features, config.DATA.NUM_CLASSESS + 1)
self.detector = yolov4(inference=True,
n_classes=config.DATA.NUM_CLASSESS)
# """
# 预训练模型
# """
# pretrained_dict = torch.load('pretrained/yolov4.pth')
# self.detector.load_state_dict(pretrained_dict)
self.yolov4loss = Yolo_loss(device=opt.device, batch=opt.batch_size)
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
def update(self, sample, *arg):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': input_image [b, 3, height, width],
'bboxes': bboxes [b, None, 4],
'labels': labels [b, None],
'path': paths}
"""
images = sample['image'].to(opt.device)
bboxes = sample['yolo4_boxes'].to(opt.device)
# bboxes = [torch.cat([box, label.unsqueeze(1)], dim=1).to(opt.device) for box, label in zip(sample['bboxes'],sample['labels'])]
out = self.detector(images)
loss, loss_xy, loss_wh, loss_obj, loss_cls, loss_l2 = self.yolov4loss(
out, bboxes)
self.avg_meters.update({'loss': loss.item()})
# 'loss_xy':loss_xy.item(), 'loss_wh':loss_wh.item(), 'loss_obj':loss_obj.item(),
# 'loss_cls':loss_cls.item(), 'loss_l2':loss_l2.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward_test(self, image):
"""给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用"""
conf_thresh = 0.001
nms_thresh = 0.45
batch_bboxes = []
batch_labels = []
batch_scores = []
box_array, confs = self.detector(image)
# [batch, num, 1, 4], num=16128
box_array = box_array[:, :, 0] #.detach().cpu().numpy()
# [batch, num, num_classes]
# confs = confs.detach().cpu().numpy()
# ipdb.set_trace()
# max_conf = np.max(confs, axis=2)
# max_id = np.argmax(confs, axis=2)
max_conf, max_id = torch.max(confs, dim=2)
b = box_array.shape[0]
for i in range(b):
argwhere = max_conf[i] > conf_thresh
l_box_array = box_array[i, argwhere, :]
l_max_conf = max_conf[i, argwhere]
l_max_id = max_id[i, argwhere]
boxes = []
labels = []
scores = []
for j in range(config.DATA.NUM_CLASSESS):
cls_argwhere = l_max_id == j
ll_box_array = l_box_array[cls_argwhere, :]
ll_max_conf = l_max_conf[cls_argwhere]
ll_max_id = l_max_id[cls_argwhere]
# keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)
keep = nms(ll_box_array, ll_max_conf, nms_thresh)
if (keep.shape[0] > 0):
ll_box_array = ll_box_array[keep]
ll_box_array = torch.clamp(ll_box_array, min=0, max=1)
ll_box_array[:, 0::2] *= opt.width
ll_box_array[:, 1::2] *= opt.height
ll_max_conf = ll_max_conf[keep]
ll_max_id = ll_max_id[keep]
boxes.extend(ll_box_array.detach().cpu().numpy().tolist())
labels.extend(ll_max_id.detach().cpu().numpy().tolist())
scores.extend(ll_max_conf.detach().cpu().numpy().tolist())
else:
pass
batch_bboxes.append(np.array(boxes))
batch_labels.append(np.array(labels).astype(np.int32))
batch_scores.append(np.array(scores))
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch, published=False):
super(Model, self).save(which_epoch, published=published)
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
kargs = {}
if 'SCALE' in config.DATA:
scale = config.DATA.SCALE
if isinstance(scale, int):
min_size = scale
max_size = int(min_size / 3 * 5)
else:
min_size, max_size = config.DATA.SCALE
kargs = {
'min_size': min_size,
'max_size': max_size,
}
kargs.update({'box_nms_thresh': config.TEST.NMS_THRESH})
# 多卡使用 SyncBN
if is_distributed():
kargs.update({'norm_layer': torch.nn.SyncBatchNorm})
# 定义backbone和Faster RCNN模型
if config.MODEL.BACKBONE is None or config.MODEL.BACKBONE.lower() in [
'res50', 'resnet50'
]:
# 默认是带fpn的resnet50
self.detector = fasterrcnn_resnet50_fpn(pretrained=False, **kargs)
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, config.DATA.NUM_CLASSESS + 1)
elif config.MODEL.BACKBONE.lower() in ['vgg16', 'vgg']:
backbone = vgg16_backbone()
self.detector = FasterRCNN(backbone,
num_classes=config.DATA.NUM_CLASSESS +
1,
**kargs)
elif config.MODEL.BACKBONE.lower() in ['res101', 'resnet101']:
# 不带FPN的resnet101
backbone = res101_backbone()
self.detector = FasterRCNN(backbone,
num_classes=config.DATA.NUM_CLASSESS +
1,
**kargs)
elif config.MODEL.BACKBONE.lower() in ['res', 'resnet']:
raise RuntimeError(
f'backbone "{config.MODEL.BACKBONE}" is ambiguous, please specify layers.'
)
else:
raise NotImplementedError(
f'no such backbone: {config.MODEL.BACKBONE}')
if opt.debug and is_first_gpu():
print_network(self.detector)
self.to(opt.device)
# 多GPU支持
if is_distributed():
self.detector = torch.nn.parallel.DistributedDataParallel(
self.detector,
find_unused_parameters=False,
device_ids=[opt.local_rank],
output_device=opt.local_rank)
# self.detector = torch.nn.parallel.DistributedDataParallel(self.detector, device_ids=[opt.local_rank], output_device=opt.local_rank)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.detector = get_net().to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *arg):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
loss = self.forward(sample)
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward(self, sample):
labels = sample['labels']
for i in range(len(labels)):
labels[i] += 1 # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
for bbox in bboxes:
bbox[:, [0, 1, 2, 3]] = bbox[:, [1, 0, 3, 2]] # yxyx
image = image.to(opt.device)
bboxes = [bbox.to(opt.device).float() for bbox in bboxes]
labels = [label.to(opt.device).float() for label in labels]
# import ipdb
# ipdb.set_trace()
loss, _, _ = self.detector(image, bboxes, labels)
return loss
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
total_loss = 0.
for i, sample in enumerate(dataloader):
utils.progress_bar(i, len(dataloader), 'Eva... ')
with torch.no_grad():
loss = self.forward(sample)
total_loss += loss.item()
logger.info(
f'Eva({data_name}) epoch {epoch}, total loss: {total_loss}.')
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, config, **kwargs):
super(Model, self).__init__(config, kwargs)
self.config = config
self.detector = SSDDetector(config).to(device=opt.device)
#####################
# Init weights
#####################
# normal_init(self.detector)
if opt.debug:
print_network(self.detector)
self.optimizer = get_optimizer(config, self.detector)
self.scheduler = get_scheduler(config, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join('checkpoints', opt.tag)
CENTER_VARIANCE = 0.1
SIZE_VARIANCE = 0.2
THRESHOLD = 0.5
self.target_transform = SSDTargetTransform(
PriorBox(config)(), CENTER_VARIANCE, SIZE_VARIANCE, THRESHOLD)
def update(self, sample, *arg):
"""
给定一个batch的图像和gt, 更新网络权重, 仅在训练时使用.
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
loss_dict = self.forward_train(sample)
loss = sum(loss for loss in loss_dict.values())
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward_train(self, sample):
labels = sample['labels']
for label in labels:
label += 1. # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
image = image * 255
sub = torch.Tensor([[123, 117, 104]]).view([1, 3, 1, 1])
image -= sub
for box in bboxes:
box[:, 0::2] /= 512 # width
box[:, 1::2] /= 512 # height
for b in range(len(labels)):
bboxes[b], labels[b] = self.target_transform(bboxes[b], labels[b])
image = image.to(opt.device)
bboxes = torch.stack(bboxes).to(opt.device)
labels = torch.stack(labels).long().to(opt.device)
targets = {'boxes': bboxes, 'labels': labels}
loss_dict = self.detector(image, targets=targets)
return loss_dict
def forward_test(self, image):
"""给定一个batch的图像, 输出预测的[bounding boxes, labels和scores], 仅在验证和测试时使用"""
conf_thresh = 0.000
batch_bboxes = []
batch_labels = []
batch_scores = []
"""
图像预处理
"""
image = image * 255
sub = torch.Tensor([[123, 117, 104]]).cuda().view([1, 3, 1, 1])
image -= sub
outputs = self.detector(image)
"""
nms阈值设置在 .box_head/inference.py
"""
for b in range(len(outputs)): #
output = outputs[b]
boxes = output['boxes']
labels = output['labels']
scores = output['scores']
labels = labels - 1
batch_bboxes.append(boxes.detach().cpu().numpy())
batch_labels.append(labels.detach().cpu().numpy())
batch_scores.append(scores.detach().cpu().numpy())
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch, published=False):
super(Model, self).save(which_epoch, published)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
# 根据YoloV2和YoloV3使用不同的配置文件
if opt.model == 'Yolo2':
cfgfile = 'configs/yolo2-voc.cfg'
elif opt.model == 'Yolo3':
cfgfile = 'configs/yolo3-coco.cfg'
# 初始化detector
self.detector = Darknet(cfgfile, device=opt.device).to(opt.device)
print_network(self.detector)
# 在--load之前加载weights文件(可选)
if opt.weights:
utils.color_print('Load Yolo weights from %s.' % opt.weights, 3)
self.detector.load_weights(opt.weights)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *args):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
loss_layers = self.detector.loss_layers
org_loss = []
image = sample['image'].to(opt.device) # target domain
target = sample['yolo_boxes'].to(opt.device)
detection_output = self.detector(image) # 在src domain上训练检测
for i, l in enumerate(loss_layers):
ol=l(detection_output[i]['x'], target)
org_loss.append(ol)
loss = sum(org_loss)
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.detector.parameters(), 10000)
self.optimizer.step()
org_loss.clear()
gc.collect()
return {}
def forward(self, image):
"""
Args:
image: [b, 3, h, w] Tensor
"""
batch_bboxes = []
batch_labels = []
batch_scores = []
if self.detector.net_name() == 'region': # region_layer
shape = (0, 0)
else:
shape = (opt.width, opt.height)
num_classes = self.detector.num_classes
outputs = self.detector(image)
outputs = get_all_boxes(outputs, shape, conf_thresh, num_classes,
device=opt.device, only_objectness=False,
validation=True)
# ipdb.set_trace()
for b in range(image.shape[0]):
preds = outputs[b]
preds = preds[preds[:, 4] > conf_thresh]
boxes = preds[:, :4]
det_conf = preds[:, 4]
cls_conf = preds[:, 5:]
b, c, h, w = image.shape
boxes = xywh_to_xyxy(boxes, w, h) # yolo的xywh转成输出的xyxy
nms_indices = nms(boxes, det_conf, nms_thresh)
if len(nms_indices) == 0:
batch_bboxes.append(np.array([[]], np.float32))
batch_labels.append(np.array([], np.int32))
batch_scores.append(np.array([], np.float32))
continue
boxes, det_conf, cls_conf = keep(nms_indices, [boxes, det_conf, cls_conf])
max_conf, labels = torch.max(cls_conf, 1)
scores = det_conf * max_conf
if not DO_FAST_EVAL: # 只保留概率最高的类，能够加快eval速度但会降低精度
repeat_boxes = boxes.repeat([num_classes, 1])
repeat_det_conf = det_conf.repeat(num_classes)
repeat_conf = cls_conf.transpose(1, 0).contiguous().view(-1)
repeat_labels = torch.linspace(0, num_classes-1, num_classes).repeat(boxes.shape[0], 1).transpose(1, 0).contiguous().view(-1)
boxes, det_conf, cls_conf, labels = keep(repeat_conf>cls_thresh, [repeat_boxes, repeat_det_conf, repeat_conf, repeat_labels])
scores = det_conf * cls_conf
batch_bboxes.append(to_numpy(boxes))
batch_labels.append(to_numpy(labels, np.int32))
batch_scores.append(to_numpy(scores))
return batch_bboxes, batch_labels, batch_scores
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)
class Model(BaseModel):
def __init__(self, opt):
super(Model, self).__init__()
self.opt = opt
self.detector = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=False)
# self.detector = FasterRCNN_VGG()
in_features = self.detector.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
self.detector.roi_heads.box_predictor = FastRCNNPredictor(
in_features, opt.num_classes + 1)
print_network(self.detector)
self.optimizer = get_optimizer(opt, self.detector)
self.scheduler = get_scheduler(opt, self.optimizer)
self.avg_meters = ExponentialMovingAverage(0.95)
self.save_dir = os.path.join(opt.checkpoint_dir, opt.tag)
def update(self, sample, *arg):
"""
Args:
sample: {'input': a Tensor [b, 3, height, width],
'bboxes': a list of bboxes [[N1 × 4], [N2 × 4], ..., [Nb × 4]],
'labels': a list of labels [[N1], [N2], ..., [Nb]],
'path': a list of paths}
"""
labels = sample['labels']
for label in labels:
label += 1. # effdet的label从1开始
image, bboxes, labels = sample['image'], sample['bboxes'], sample[
'labels']
if len(bboxes[0]) == 0: # 没有bbox，不更新参数
return {}
image = image.to(opt.device)
bboxes = [bbox.to(opt.device).float() for bbox in bboxes]
labels = [label.to(opt.device).float() for label in labels]
image = list(im for im in image)
b = len(bboxes)
target = [{
'boxes': bboxes[i],
'labels': labels[i].long()
} for i in range(b)]
"""
target['boxes'] = boxes
target['labels'] = labels
# target['masks'] = None
target['image_id'] = torch.tensor([index])
target['area'] = area
target['iscrowd'] = iscrowd
"""
loss_dict = self.detector(image, target)
loss = sum(l for l in loss_dict.values())
self.avg_meters.update({'loss': loss.item()})
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return {}
def forward(self, image): # test
conf_thresh = 0.5
image = list(im for im in image)
batch_bboxes = []
batch_labels = []
batch_scores = []
with torch.no_grad():
outputs = self.detector(image)
for b in range(len(outputs)): #
output = outputs[b]
boxes = output['boxes']
labels = output['labels']
scores = output['scores']
boxes = boxes[scores > conf_thresh]
labels = labels[scores > conf_thresh]
labels = labels.detach().cpu().numpy()
# for i in range(len(labels)):
# labels[i] = coco_90_to_80_classes(labels[i])
labels = labels - 1
scores = scores[scores > conf_thresh]
batch_bboxes.append(boxes.detach().cpu().numpy())
batch_labels.append(labels)
batch_scores.append(scores.detach().cpu().numpy())
return batch_bboxes, batch_labels, batch_scores
def inference(self, x, progress_idx=None):
raise NotImplementedError
def evaluate(self, dataloader, epoch, writer, logger, data_name='val'):
return self.eval_mAP(dataloader, epoch, writer, logger, data_name)
def load(self, ckpt_path):
return super(Model, self).load(ckpt_path)
def save(self, which_epoch):
super(Model, self).save(which_epoch)