diff --git a/yolov6/assigners/anchor_generator.py b/yolov6/assigners/anchor_generator.py index 21968487..c8276418 100644 --- a/yolov6/assigners/anchor_generator.py +++ b/yolov6/assigners/anchor_generator.py @@ -1,5 +1,7 @@ import torch +from yolov6.utils.general import check_version +torch_1_10_plus = check_version(torch.__version__, minimum='1.10.0') def generate_anchors(feats, fpn_strides, grid_cell_size=5.0, grid_cell_offset=0.5, device='cpu', is_eval=False, mode='af'): '''Generate anchors from features.''' @@ -13,7 +15,7 @@ def generate_anchors(feats, fpn_strides, grid_cell_size=5.0, grid_cell_offset=0. _, _, h, w = feats[i].shape shift_x = torch.arange(end=w, device=device) + grid_cell_offset shift_y = torch.arange(end=h, device=device) + grid_cell_offset - shift_y, shift_x = torch.meshgrid(shift_y, shift_x, indexing='ij') + shift_y, shift_x = torch.meshgrid(shift_y, shift_x, indexing='ij') if torch_1_10_plus else torch.meshgrid(shift_y, shift_x) anchor_point = torch.stack( [shift_x, shift_y], axis=-1).to(torch.float) if mode == 'af': # anchor-free @@ -35,7 +37,7 @@ def generate_anchors(feats, fpn_strides, grid_cell_size=5.0, grid_cell_offset=0. cell_half_size = grid_cell_size * stride * 0.5 shift_x = (torch.arange(end=w, device=device) + grid_cell_offset) * stride shift_y = (torch.arange(end=h, device=device) + grid_cell_offset) * stride - shift_y, shift_x = torch.meshgrid(shift_y, shift_x, indexing='ij') + shift_y, shift_x = torch.meshgrid(shift_y, shift_x, indexing='ij') if torch_1_10_plus else torch.meshgrid(shift_y, shift_x) anchor = torch.stack( [ shift_x - cell_half_size, shift_y - cell_half_size, diff --git a/yolov6/data/data_augment.py b/yolov6/data/data_augment.py index 76c4b519..45df88e6 100644 --- a/yolov6/data/data_augment.py +++ b/yolov6/data/data_augment.py @@ -78,7 +78,10 @@ def random_affine(img, labels=(), degrees=10, translate=.1, scale=.1, shear=10, new_shape=(640, 640)): '''Applies Random affine transformation.''' n = len(labels) - height, width = new_shape + if isinstance(new_shape, int): + height = width = new_shape + else: + height, width = new_shape M, s = get_transform_matrix(img.shape[:2], (height, width), degrees, scale, shear, translate) if (M != np.eye(3)).any(): # image changed diff --git a/yolov6/utils/general.py b/yolov6/utils/general.py index ec88a434..cb4418cd 100644 --- a/yolov6/utils/general.py +++ b/yolov6/utils/general.py @@ -5,6 +5,7 @@ import math import torch import requests +import pkg_resources as pkg from pathlib import Path from yolov6.utils.events import LOGGER @@ -114,3 +115,13 @@ def check_img_size(imgsz, s=32, floor=0): if new_size != imgsz: LOGGER.warning(f'--img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}') return new_size + + +def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False, verbose=False): + # Check whether the package's version is match the required version. + current, minimum = (pkg.parse_version(x) for x in (current, minimum)) + result = (current == minimum) if pinned else (current >= minimum) # bool + if hard: + info = f'⚠️ {name}{minimum} is required by YOLOv6, but {name}{current} is currently installed' + assert result, info # assert minimum version requirement + return result