Browse Source

update mosaic border

5.0
Glenn Jocher 4 years ago
parent
commit
93a6765806
1 changed files with 17 additions and 13 deletions
  1. +17
    -13
      utils/datasets.py

+ 17
- 13
utils/datasets.py View File





class LoadImages: # for inference class LoadImages: # for inference
def __init__(self, path, img_size=416):
def __init__(self, path, img_size=640):
path = str(Path(path)) # os-agnostic path = str(Path(path)) # os-agnostic
files = [] files = []
if os.path.isdir(path): if os.path.isdir(path):




class LoadWebcam: # for inference class LoadWebcam: # for inference
def __init__(self, pipe=0, img_size=416):
def __init__(self, pipe=0, img_size=640):
self.img_size = img_size self.img_size = img_size


if pipe == '0': if pipe == '0':




class LoadStreams: # multiple IP or RTSP cameras class LoadStreams: # multiple IP or RTSP cameras
def __init__(self, sources='streams.txt', img_size=416):
def __init__(self, sources='streams.txt', img_size=640):
self.mode = 'images' self.mode = 'images'
self.img_size = img_size self.img_size = img_size






class LoadImagesAndLabels(Dataset): # for training/testing class LoadImagesAndLabels(Dataset): # for training/testing
def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_images=False, single_cls=False, stride=32, pad=0.0): cache_images=False, single_cls=False, stride=32, pad=0.0):
try: try:
path = str(Path(path)) # os-agnostic path = str(Path(path)) # os-agnostic
self.image_weights = image_weights self.image_weights = image_weights
self.rect = False if image_weights else rect self.rect = False if image_weights else rect
self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training) self.mosaic = self.augment and not self.rect # load 4 images at a time into a mosaic (only during training)
self.mosaic_border = None
self.stride = stride



# Define labels # Define labels
self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt') self.label_files = [x.replace('images', 'labels').replace(os.path.splitext(x)[-1], '.txt')


labels4 = [] labels4 = []
s = self.img_size s = self.img_size
xc, yc = [int(random.uniform(s * 0.5, s * 1.5)) for _ in range(2)] # mosaic center x, y
border = [-s // 2, -s // 2] # self.mosaic_border
yc, xc = [int(random.uniform(-x, 2 * s + x)) for x in border] # mosaic center x, y
indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)] # 3 additional image indices indices = [index] + [random.randint(0, len(self.labels) - 1) for _ in range(3)] # 3 additional image indices
for i, index in enumerate(indices): for i, index in enumerate(indices):
# Load image # Load image
translate=self.hyp['translate'], translate=self.hyp['translate'],
scale=self.hyp['scale'], scale=self.hyp['scale'],
shear=self.hyp['shear'], shear=self.hyp['shear'],
border=-s // 2) # border to remove
border=border) # border to remove


return img4, labels4 return img4, labels4




def letterbox(img, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
# Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
shape = img.shape[:2] # current shape [height, width] shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int): if isinstance(new_shape, int):
return img, ratio, (dw, dh) return img, ratio, (dw, dh)




def random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, border=0):
def random_affine(img, targets=(), degrees=10, translate=.1, scale=.1, shear=10, border=(0, 0)):
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10)) # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
# https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4 # https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4
# targets = [cls, xyxy] # targets = [cls, xyxy]


height = img.shape[0] + border * 2
width = img.shape[1] + border * 2
height = img.shape[0] + border[0] * 2 # shape(h,w,c)
width = img.shape[1] + border[1] * 2


# Rotation and Scale # Rotation and Scale
R = np.eye(3) R = np.eye(3)


# Translation # Translation
T = np.eye(3) T = np.eye(3)
T[0, 2] = random.uniform(-translate, translate) * img.shape[0] + border # x translation (pixels)
T[1, 2] = random.uniform(-translate, translate) * img.shape[1] + border # y translation (pixels)
T[0, 2] = random.uniform(-translate, translate) * img.shape[1] + border[1] # x translation (pixels)
T[1, 2] = random.uniform(-translate, translate) * img.shape[0] + border[0] # y translation (pixels)


# Shear # Shear
S = np.eye(3) S = np.eye(3)


# Combined rotation matrix # Combined rotation matrix
M = S @ T @ R # ORDER IS IMPORTANT HERE!! M = S @ T @ R # ORDER IS IMPORTANT HERE!!
if (border != 0) or (M != np.eye(3)).any(): # image changed
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114)) img = cv2.warpAffine(img, M[:2], dsize=(width, height), flags=cv2.INTER_LINEAR, borderValue=(114, 114, 114))


# Transform label coordinates # Transform label coordinates

Loading…
Cancel
Save