def __init__(self, backbone=None, neck=None, head=None, pretrained=None):
super(YOLOv4Detector, self).__init__()
self.backbone = build_from_dict(backbone, BACKBONES)
if neck is not None:
self.neck = build_from_dict(neck, NECKS)
if head is not None:
self.head = build_from_dict(head, HEADS)
self.init_weights(pretrained)
def __init__(self,
label_smooth=True,
conf_balances=[0.4, 1, 4],
deta=0.01,
anchors=[[12, 16], [19, 36], [40, 28], [36, 75], [76, 55],
[72, 146], [142, 110], [192, 243], [459, 401]],
anchor_masks=[[6, 7, 8], [3, 4, 5], [0, 1, 2]],
downsample_ratios=[32, 16, 8],
bbox_weight=False,
yolo_loss_type=None,
in_channels=[128, 256, 512],
num_classes=80,
nms_type='nms',
nms_thr=.5,
ignore_thre=.3,
anchor_t=4,
bbox_loss=dict(type='IOU_Loss', iou_type='CIOU'),
confidence_loss=dict(type='Conf_Loss'),
class_loss=dict(type='Class_Loss'),
norm_type='BN',
num_groups=None):
super(BaseHead, self).__init__()
assert yolo_loss_type in (
None, 'yolov4',
'yolov5'), 'use_yolo_loss just support [None,yolov4,yolov5]'
self.yolo_loss_type = yolo_loss_type
self.conf_balances = conf_balances
self.num_classes = num_classes
assert nms_type in [
'nms', 'soft_nms'
], 'nms type only support [nms,soft_nms],other has not implementation'
self.nms_thr = nms_thr
self.nms_type = nms_type
self.out_channels = []
self.in_channels = in_channels
self.norm_type = norm_type
self.num_groups = num_groups
self.base_num = 5
self.bbox_weight = bbox_weight
self.anchor_t = anchor_t
for mask in anchor_masks:
self.out_channels.append(len(mask) * (self.base_num + num_classes))
self.anchors = anchors
self.anchor_masks = anchor_masks
self.downsample_ratios = downsample_ratios
if label_smooth is None or not isinstance(label_smooth, bool):
label_smooth = False
self.label_smooth = label_smooth
self.deta = deta
self.ignore_thre = ignore_thre
self.bbox_loss = build_from_dict(bbox_loss, LOSS)
self.confidence_loss = build_from_dict(confidence_loss, LOSS)
self.class_loss = build_from_dict(class_loss, LOSS)
def __init__(self,
batch=8,
subdivisions=4,
epochs=100,
burn_in=1000,
steps=[400000, 450000]):
_model = build_from_dict(model, DETECTORS)
self.model = DataParallel(_model.cuda(), device_ids=[0])
self.train_dataset = build_from_dict(data_cfg['train'], DATASET)
self.val_dataset = build_from_dict(data_cfg['val'], DATASET)
self.burn_in = burn_in
self.steps = steps
self.epochs = epochs
self.batch = batch
self.subdivisions = subdivisions
self.train_size = len(self.train_dataset)
self.val_size = len(self.val_dataset)
self.train_loader = DataLoader(self.train_dataset,
batch_size=batch // subdivisions,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True,
collate_fn=self.collate)
self.val_loader = DataLoader(self.val_dataset,
batch_size=batch // subdivisions,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True,
collate_fn=self.collate)
self.optimizer = optim.Adam(
self.model.parameters(),
lr=0.001 / batch,
betas=(0.9, 0.999),
eps=1e-08,
)
self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer,
self.burnin_schedule)
def register_checkpoint_hook(self, checkpoint_config):
if checkpoint_config is None:
return
if isinstance(checkpoint_config, dict):
checkpoint_config.setdefault('type', 'CheckpointHook')
hook = build_from_dict(checkpoint_config, HOOKS)
else:
hook = checkpoint_config
self.register_hook(hook)
def register_lr_hook(self, lr_config):
if isinstance(lr_config, dict):
assert 'policy' in lr_config
hook_type = lr_config.pop('policy').title() + 'LrUpdaterHook'
lr_config['type'] = hook_type
hook = build_from_dict(lr_config, HOOKS)
else:
hook = lr_config
self.register_hook(hook)
def register_optimizer_hook(self, optimizer_config):
if optimizer_config is None:
return
if isinstance(optimizer_config, dict):
optimizer_config.setdefault('type', 'OptimizerHook')
hook = build_from_dict(optimizer_config, HOOKS)
else:
hook = optimizer_config
self.register_hook(hook)
def collate(self, batch):
if 'multi_scale' in data_cfg.keys() and len(
data_cfg['multi_scale']) > 0:
multi_scale = data_cfg['multi_scale']
if isinstance(multi_scale, dict) and 'type' in multi_scale.keys():
randomShape = build_from_dict(multi_scale, TRANSFORMS)
batch = randomShape(batch)
collate = default_collate(batch)
return collate
def train_detector(model,
dataset,
cfg,
validate=False,
timestamp=None,
meta=None):
logger = Logging.getLogger()
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
data_loaders = [build_dataloader(ds, data=cfg.data) for ds in dataset]
if torch.cuda.is_available():
model = model.cuda(cfg.gpu_ids[0])
model.device = cfg.gpu_ids[0]
if torch.cuda.device_count() > 1:
model = DataParallel(model, device_ids=cfg.gpu_ids)
else:
model.device = 'cpu'
# build runner
optimizer = cfg.optimizer
if 'ema' in cfg:
ema = cfg.ema
else:
ema = None
runner = Runner(model,
batch_processor,
optimizer,
cfg.work_dir,
logger=logger,
meta=meta,
ema=ema)
# an ugly walkaround to make the .log and .log.json filenames the same
runner.timestamp = timestamp
# register eval hooks 需要放在日志前面，不然打印不出日志。
if validate:
cfg.data.val.train = False
val_dataset = build_from_dict(cfg.data.val, DATASET)
val_dataloader = build_dataloader(val_dataset,
shuffle=False,
data=cfg.data)
eval_cfg = cfg.get('evaluation', {})
from yolodet.models.hooks.eval_hook import EvalHook
runner.register_hook(EvalHook(val_dataloader, **eval_cfg))
# register hooks
# runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,cfg.checkpoint_config)
runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
cfg.checkpoint_config, cfg.log_config)
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
def __init__(self, transforms):
assert isinstance(transforms, collections.abc.Sequence)
self.transforms = []
for transform in transforms:
if isinstance(transform, dict):
transform = build_from_dict(transform, TRANSFORMS)
self.transforms.append(transform)
elif callable(transform):
self.transforms.append(transform)
else:
raise TypeError('transform must be callable or a dict')
def __init__(self,
depth_multiple,
width_multiple,
backbone=None,
neck=None,
head=None,
pretrained=None):
super(YOLOv5Detector, self).__init__()
self.depth_multiple = depth_multiple
self.width_multiple = width_multiple
backbone['depth_multiple'] = depth_multiple
backbone['width_multiple'] = width_multiple
self.backbone = build_from_dict(backbone, BACKBONES)
if neck is not None:
neck['depth_multiple'] = depth_multiple
neck['width_multiple'] = width_multiple
self.neck = build_from_dict(neck, NECKS)
if head is not None:
head['depth_multiple'] = depth_multiple
head['width_multiple'] = width_multiple
self.head = build_from_dict(head, HEADS)
self.init_weights(pretrained)
def __init__(self,
aware_Loss=dict(type='IOU_Aware_Loss'),
scale_x_y=1.05,
iou_aware=True,
iou_aware_factor=0.4,
coord_conv=True,
xywh_loss=True,
**kwargs):
super(PPHead, self).__init__(**kwargs)
self.coord_conv = coord_conv
self.scale_x_y = scale_x_y
self.iou_aware = iou_aware
self.iou_aware_factor = iou_aware_factor
self.xywh_loss = xywh_loss
out_channels = []
self.base_num = 5
if iou_aware:
self.base_num = 6
for mask in self.anchor_masks:
out_channels.append(len(mask) * (self.base_num + self.num_classes))
self.y1 = Y(self.in_channels[-1],
out_channels[0],
norm_type=self.norm_type,
num_groups=self.num_groups,
coord_conv=self.coord_conv)
self.y2 = Y(self.in_channels[-2],
out_channels[1],
norm_type=self.norm_type,
num_groups=self.num_groups,
coord_conv=self.coord_conv)
self.y3 = Y(self.in_channels[-3],
out_channels[2],
norm_type=self.norm_type,
num_groups=self.num_groups,
coord_conv=self.coord_conv)
if aware_Loss is not None:
self.aware_Loss = build_from_dict(aware_Loss, LOSS)
import cv2
import os
import matplotlib.pyplot as plt
def draw_box(img, bboxes,class_name,gt_class,gt_score):
# for b in bboxes:
# img = cv2.rectangle(img, (b[0], b[1]), (b[2], b[3]), (0, 255, 0), 1)
for idx,ind in enumerate(gt_class):
if ind>-1:
b = bboxes[idx]
img = cv2.rectangle(img, (b[0], b[1]), (b[2], b[3]), (0, 255, 0), 1)
font = cv2.FONT_HERSHEY_SIMPLEX
img = cv2.putText(img,class_name[int(ind)]+'|'+str(gt_score[idx]),(b[0], b[1]), font, 0.3, (0, 0, 255), 1)
return img
dataset = build_from_dict(dataset_test.data['train'],DATASET)
for i in range(30):
result = dataset.__getitem__(i)
img = result['img']
gt_bboxes = result['gt_bboxes']
gt_score = result['gt_score']
gt_bboxes[:, [1, 3]] *= img.shape[0] # height
gt_bboxes[:, [0, 2]] *= img.shape[1] # width
img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2BGR)
a = draw_box(img.copy(), gt_bboxes.astype(np.int32),dataset.CLASSES,result['gt_class'],gt_score)
for b in gt_bboxes.astype(np.int32):
if b[0] !=0 or b[1] !=0 or b[2] !=0 or b[3] !=0:
print(b)
cv2.imwrite(os.path.join('/disk2/project/test/v2.0/yolov5/dataset/123/test',str(i)+'.jpg'),a)
def register_logger_hooks(self, log_config):
log_interval = log_config['interval']
for info in log_config['hooks']:
logger_hook = build_from_dict(
info, HOOKS, default_args=dict(interval=log_interval))
self.register_hook(logger_hook)
for idx,bbox in enumerate(bboxes):
# x1, y1, x2, y2 = bbox
cx, cy, _w, _h = bbox
x, y = cx - _w / 2, cy - _h / 2
# x1,y1,x2,y2 = int(x1*w),int(y1*h),int(x2*w),int(y2*h)
x1,y1,x2,y2 = x,y, x+_w,y+_h
x1, y1, x2, y2 = int(x1 * w), int(y1 * h), int(x2 * w), int(y2 * h)
# img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1)
img = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1)
return img
file = '/disk2/project/pytorch-YOLOv4/cfg/dataset_test.py'
cfg = Config.fromfile(file)
dataset = build_from_dict(cfg.data.train,DATASET)
dataloader = build_dataloader(dataset,data=cfg.data)
for i, data_batch in enumerate(dataloader):
if i>30:
break
for idx,data in enumerate(data_batch['img']):
gt = data_batch['gt_bboxes'][idx]
gt_xywh = xyxy2xywh(gt) # x,y ,w, h
n_gt = (gt.sum(dim=-1) > 0).sum(dim=-1)
n = int(n_gt)
if n == 0:
continue
gt = gt[:n].cpu().numpy()
gt_xywh = gt_xywh[:n].cpu().numpy()
return img
def filter_pipelines(cfg_dataset, filter=[]):
if len(filter) == 0:
print(f'None type for filter\n')
return cfg_dataset
pipelines = []
for pipeline in cfg_dataset:
if pipeline['type'] not in filter:
pipelines.append(pipeline)
return pipelines
yolov5_coco_100e.data['train']['pipeline'] = filter_pipelines(yolov5_coco_100e.data['train']['pipeline'],\
['Normalize','ImageToTensor'])
dataset = build_from_dict(yolov5_coco_100e.data['train'], DATASET)
for i in range(30):
result = dataset.__getitem__(i)
img = result['img']
gt_bboxes = result['gt_bboxes']
gt_score = result['gt_score']
gt_bboxes[:, [1, 3]] *= img.shape[0] # height
gt_bboxes[:, [0, 2]] *= img.shape[1] # width
img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2BGR)
a = draw_box(img.copy(), gt_bboxes.astype(np.int32), dataset.CLASSES,
result['gt_class'], gt_score)
for b in gt_bboxes.astype(np.int32):
if b[0] != 0 or b[1] != 0 or b[2] != 0 or b[3] != 0:
print(b)
cv2.imwrite(
parser.add_argument('--verbose',
action='store_true',
help='report mAP by class')
parser.add_argument('--half',
action='store_true',
help='fp16 half precision')
opt = parser.parse_args()
print(opt)
# config = '/disk2/project/mmdetection/mount/pytorch-YOLOv4/cfg/yolov5_coco_gpu.py'
# checkpoint = '/disk2/project/pytorch-YOLOv4/work_dirs/yolov5-l_epoch_24.pth'
cfg = Config.fromfile(opt.config)
cfg.data.val.train = False
val_dataset = build_from_dict(cfg.data.val, DATASET)
val_dataloader = build_dataloader(val_dataset,
data=cfg.data,
shuffle=False)
device = select_device(opt.device)
# model = init_detector(opt.config, checkpoint=opt.checkpoint, device=device)
model = init_detector(opt.config, checkpoint=opt.checkpoint, device=device)
result = single_gpu_test(model,
val_dataloader,
half=opt.half,
conf_thres=opt.conf_thres,
iou_thres=opt.iou_thres,
merge=opt.merge,
save_json=opt.save_json,
augment=opt.augment,
verbose=opt.verbose,
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
# set cudnn_benchmark
if cfg.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
# update configs according to CLI args
if args.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
if args.device is not None:
cfg.device = args.device
else:
cfg.device = None
device = select_device(cfg.device)
if args.autoscale_lr:
# apply the linear scaling rule (https://arxiv.org/abs/1706.02677)
cfg.optimizer['lr'] = cfg.optimizer['lr'] * len(cfg.gpu_ids) / 8
# create work_dir
file_utils.mkdir_or_exist(osp.abspath(cfg.work_dir))
# init the logger before other steps
# timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
# log_file = osp.join(cfg.work_dir, '{}.log'.format(timestamp))
logger = Logging.getLogger()
# init the meta dict to record some important information such as
# environment info and seed, which will be logged
meta = dict()
# log env info
env_info_dict = collect_env()
env_info = '\n'.join([('{}: {}'.format(k, v))
for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
meta['env_info'] = env_info
meta['batch_size'] = cfg.data.batch_size
meta['subdivisions'] = cfg.data.subdivisions
meta['multi_scale'] = args.multi_scale
# log some basic info
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}, deterministic: {}'.format(
args.seed, args.deterministic))
set_random_seed(args.seed, deterministic=args.deterministic)
cfg.seed = args.seed
meta['seed'] = args.seed
model = build_from_dict(cfg.model, DETECTORS)
model = model.cuda(device)
# model.device = device
if device.type != 'cpu' and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.device = device
datasets = [build_from_dict(cfg.data.train, DATASET)]
if len(cfg.workflow) == 2:
val_dataset = copy.deepcopy(cfg.data.val)
val_dataset.pipeline = cfg.data.train.pipeline
datasets.append(build_from_dict(val_dataset, DATASET))
if cfg.checkpoint_config is not None:
# save mmdet version, config file content and class names in
# checkpoints as meta data
cfg.checkpoint_config.meta = dict(config=cfg.text,
CLASSES=datasets[0].CLASSES)
# add an attribute for visualization convenience
model.CLASSES = datasets[0].CLASSES
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
train_detector(model,
datasets,
cfg,
validate=args.validate,
timestamp=timestamp,
meta=meta)
