def augment_transform()

in sagemaker/src/word_and_line_segmentation.py [0:0]

• 26 lines of code
• 3 McCabe index (conditional complexity)

def augment_transform(image, bbox, text, random_remove_box=0.15):    
    image_h, image_w = image.shape[-2:]

    ty = random.uniform(-random_y_translation, random_y_translation)
    tx = random.uniform(-random_x_translation, random_x_translation)
    
    aug_bbs = []
    for i in range(bbox.shape[0]):
        x1, y1, x2, y2 = bbox[i]
        x1, y1, x2, y2 = x1 * image_w, y1 * image_h, x2 * image_w, y2 * image_h
        aug_bbs.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2))
    
    bbs = BoundingBoxesOnImage(aug_bbs, shape=image.shape)
    seq = iaa.Sequential([
        iaa.Affine(
            cval=255,
            translate_px={"x": int(tx*image.shape[1]), "y": int(ty*image.shape[0])},
            scale=(0.9, 1.1),
        ),
        iaa.Crop(percent=(0., 0., .3, 0.))
    ])
    image_aug, bbs_aug = seq(image=image, bounding_boxes=bbs)
    
    bbox_out = []
    for bb_aug in bbs_aug:
        x1, y1, x2, y2 = bb_aug.x1, bb_aug.y1, bb_aug.x2, bb_aug.y2
        x1, y1, x2, y2 = x1 / image_w, y1 / image_h, x2 / image_w, y2 / image_h
        bbox_out.append([x1, y1, x2, y2])
        
    bbox_out = np.array(bbox_out)
    return transform(image_aug, bbox_out, text)